admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-06 01:27:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

small impr

Browse Source
This commit is contained in:
Pepijn
2025-08-28 22:45:23 +02:00
parent 7dce022a05
commit 04d55e4670
2 changed files with 75 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/rlearn/modeling_rlearn.py
+70 -6
View File
@@ -76,6 +76,7 @@ Notes
from __future__ import annotations
import math
import time
from itertools import chain
import torch
@@ -92,7 +93,7 @@ try:
except ImportError as e:
raise ImportError(
"ReWiND dependencies not installed. Please install: "
"pip install x-transformers hl-gauss-pytorch einx einops"
"pip install x-transformers hl-gauss-pytorch einx einops x_mlps_pytorc"
) from e
from lerobot.constants import OBS_IMAGE, OBS_IMAGES, OBS_LANGUAGE, REWARD
@@ -123,12 +124,16 @@ class RLearNPolicy(PreTrainedPolicy):
from sentence_transformers import SentenceTransformer
# Load DINOv2 (base) vision encoder with its processor
self.vision_processor = AutoImageProcessor.from_pretrained(config.vision_model_name)
self.vision_processor = AutoImageProcessor.from_pretrained(config.vision_model_name, use_fast=True)
self.vision_encoder = AutoModel.from_pretrained(config.vision_model_name)
# Load sentence-transformers text encoder
self.text_encoder = SentenceTransformer(config.text_model_name)
# Move text encoder to same device as vision encoder (GPU if available)
if torch.cuda.is_available():
self.text_encoder = self.text_encoder.to('cuda')
# DINOv2-base has 768 hidden size, all-MiniLM-L12-v2 has 384
self.vision_hidden = 768 # DINOv2-base
self.text_hidden = 384 # all-MiniLM-L12-v2
@@ -296,10 +301,12 @@ class RLearNPolicy(PreTrainedPolicy):
Returns:
(B, T, D_vision)
"""
start_time = time.time()
B, T, C, H, W = frames.shape
flat = rearrange(frames, 'b t c h w -> (b t) c h w')
# Process with DINOv2
preprocess_start = time.time()
images_list = []
for i in range(B * T):
img = flat[i].permute(1, 2, 0) # CHW -> HWC
@@ -308,14 +315,29 @@ class RLearNPolicy(PreTrainedPolicy):
else:
img = (img.clamp(0, 1) * 255).to(torch.uint8).cpu().numpy()
images_list.append(img)
preprocess_time = time.time() - preprocess_start
processor_start = time.time()
processed = self.vision_processor(images=images_list, return_tensors="pt")
pixel_values = processed["pixel_values"].to(next(self.vision_encoder.parameters()).device)
processor_time = time.time() - processor_start
encoder_start = time.time()
vision_outputs = self.vision_encoder(pixel_values)
encoder_time = time.time() - encoder_start
# Extract CLS tokens
cls_tokens = vision_outputs.last_hidden_state[:, 0] # (BT, D_vision)
return rearrange(cls_tokens, '(b t) d -> b t d', b=B, t=T)
result = rearrange(cls_tokens, '(b t) d -> b t d', b=B, t=T)
total_time = time.time() - start_time
print(f"🎬 Video encoding timing (B={B}, T={T}):")
print(f" - Preprocess: {preprocess_time:.3f}s")
print(f" - Processor: {processor_time:.3f}s")
print(f" - DINOv2: {encoder_time:.3f}s")
print(f" - Total: {total_time:.3f}s")
return result
def _mask_from_lens(self, lens: Tensor) -> Tensor:
"""Create mask from sequence lengths."""
@@ -332,10 +354,13 @@ class RLearNPolicy(PreTrainedPolicy):
Note: Progress labels (0 to 1) are generated automatically for each episode.
No REWARD key is needed in the batch.
"""
forward_start = time.time()
batch = self.normalize_inputs(batch)
batch = self.normalize_targets(batch)
# Extract frames and form (B, T, C, H, W)
data_prep_start = time.time()
frames = extract_visual_sequence(batch, target_seq_len=self.config.max_seq_len)
B, T, C, H, W = frames.shape
device = next(self.parameters()).device
@@ -366,21 +391,36 @@ class RLearNPolicy(PreTrainedPolicy):
commands = [""] * B
elif not isinstance(commands, list):
commands = [str(commands)] * B
data_prep_time = time.time() - data_prep_start
# Process video frames through DINOv2
video_embeds = self._encode_video_frames(frames) # (B, T_eff, D_vision)
video_embeds = self._encode_video_frames(frames) # (B, T_eff, D_vision) - timing inside
# Language embeddings
lang_start = time.time()
print(f"🔍 Text encoder device: {next(self.text_encoder.parameters()).device if hasattr(self.text_encoder, 'parameters') else 'Unknown'}")
print(f"🔍 Target device: {device}")
print(f"🔍 Commands: {len(commands)} items, first: '{commands[0][:50]}...'")
lang_embeds = self.text_encoder.encode(
commands,
output_value='token_embeddings',
convert_to_tensor=True,
device=device
)
encode_time = time.time() - lang_start
pad_start = time.time()
lang_embeds = pad_sequence(lang_embeds, batch_first=True).to(device)
lens = torch.tensor([le.shape[0] for le in lang_embeds], device=device)
mask = self._mask_from_lens(lens)
pad_time = time.time() - pad_start
lang_time = time.time() - lang_start
print(f"🗣️ Language breakdown: encode={encode_time:.3f}s, pad={pad_time:.3f}s, total={lang_time:.3f}s")
# Token preparation
token_prep_start = time.time()
# Register tokens
register_tokens = repeat(self.register_tokens, 'n d -> b n d', b=B)
@@ -398,15 +438,20 @@ class RLearNPolicy(PreTrainedPolicy):
# Extend mask for register and video tokens
mask = F.pad(mask, (0, register_tokens.shape[1] + video_tokens.shape[1]), value=True)
token_prep_time = time.time() - token_prep_start
# Forward through x_transformers Decoder
transformer_start = time.time()
attended = self.decoder(tokens, mask=mask)
transformer_time = time.time() - transformer_start
# Unpack and get video token features
unpack_start = time.time()
_, _, attended_video_tokens = unpack(attended, lang_video_packed_shape, 'b * d')
# MLP predictor
video_frame_embeds = self.mlp_predictor(attended_video_tokens)
unpack_time = time.time() - unpack_start
# Generate progress labels on-the-fly (ReWiND approach)
# IMPORTANT: Progress should be 0-1 across the ENTIRE EPISODE, not just the temporal window
@@ -441,7 +486,7 @@ class RLearNPolicy(PreTrainedPolicy):
progress_values.append(progress)
# Create progress tensor for the current frame (last in temporal sequence)
current_progress = torch.tensor(progress_values, device=values.device, dtype=values.dtype)
current_progress = torch.tensor(progress_values, device=video_frame_embeds.device, dtype=video_frame_embeds.dtype)
# Now calculate progress for ALL frames in the temporal window
# The observation_delta_indices tell us which frames we're looking at
@@ -472,7 +517,7 @@ class RLearNPolicy(PreTrainedPolicy):
frame_progress.append(prog)
all_progress.append(
torch.tensor(frame_progress, device=values.device, dtype=values.dtype)
torch.tensor(frame_progress, device=video_frame_embeds.device, dtype=video_frame_embeds.dtype)
)
# Stack to get (B, T) tensor where T is the temporal sequence length
@@ -495,6 +540,7 @@ class RLearNPolicy(PreTrainedPolicy):
return rewards.mean() * 0.0, {"rewards_mean": rewards.mean().item()}
# Calculate loss using HLGauss or categorical
loss_start = time.time()
if self.categorical_rewards:
# Categorical cross-entropy loss
assert target.dtype in (torch.long, torch.int), "Categorical rewards require integer targets"
@@ -509,6 +555,7 @@ class RLearNPolicy(PreTrainedPolicy):
# Create video mask for variable length support
video_mask = torch.ones(B, T_eff, dtype=torch.bool, device=device)
loss = self.hl_gauss_layer(video_frame_embeds, target[:, :T_eff], mask=video_mask)
loss_time = time.time() - loss_start
# Optional: Mismatched video-language pairs loss
L_mismatch = torch.zeros((), device=device)
@@ -547,12 +594,29 @@ class RLearNPolicy(PreTrainedPolicy):
# Total loss
total_loss = loss + L_mismatch
# Calculate and print timing summary
total_forward_time = time.time() - forward_start
print(f"\n⏱️ RLearN Forward Pass Timing (B={B}, T_eff={T_eff}):")
print(f" 📊 Data prep: {data_prep_time:.3f}s ({data_prep_time/total_forward_time*100:.1f}%)")
print(f" 🗣️ Language: {lang_time:.3f}s ({lang_time/total_forward_time*100:.1f}%)")
print(f" 🔧 Token prep: {token_prep_time:.3f}s ({token_prep_time/total_forward_time*100:.1f}%)")
print(f" 🤖 Transformer: {transformer_time:.3f}s ({transformer_time/total_forward_time*100:.1f}%)")
print(f" 📦 Unpack+MLP: {unpack_time:.3f}s ({unpack_time/total_forward_time*100:.1f}%)")
print(f" 🎯 Loss calc: {loss_time:.3f}s ({loss_time/total_forward_time*100:.1f}%)")
print(f" 🏁 Total: {total_forward_time:.3f}s")
# Log individual loss components
loss_dict.update({
"loss": total_loss.item(),
"loss_main": loss.item(),
"loss_mismatch": L_mismatch.item(),
# Add timing metrics to loss dict for logging
"timing/total_forward": total_forward_time,
"timing/data_prep": data_prep_time,
"timing/language": lang_time,
"timing/transformer": transformer_time,
})
return total_loss, loss_dict
src/lerobot/policies/rlearn/rlearn_plan.md
+5 -6
View File
@@ -77,17 +77,16 @@ _ HOWTO100M: https://www.di.ens.fr/willow/research/howto100m/
- Only rewind loss [x]
- Exactly similar to: https://github.com/lucidrains/rewind-reward-pytorch/blob/main/rewind_reward_pytorch/rewind_reward.py#L11 [x]
- Try DINO v2 as encoder Base 86 M: with https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2 [x]
- Test only rewind loss (evaluate) []
- Only vlc loss then eval []
- Vlc + Rewind loss then eval []
- Cleanup code []
- Convert python -m lerobot.datasets.v21.convert_dataset_v20_to_v21 --repo-id=IPEC-COMMUNITY/bc_z_lerobot and train on 1 percent
- Test rewind (evaluate) []
- Cleanup code? []
- Convert python -m lerobot.datasets.v21.convert_dataset_v20_to_v21 --repo-id=IPEC-COMMUNITY/bc_z_lerobot and train on 1 percent
-----------------
- Then on 10 percent
- Ablation dino v2 vs dino v3 base 86 M
- Add more artificial text to dataset generated by vlm (google gemini) []
- See google gemini vlm caption [] https://gemini.google.com/app/7e332ffaf32580f2
- Multiple captions per video, creat method to generate as much data as possible etc [] https://arxiv.org/abs/2508.13446, https://arxiv.org/pdf/2412.04453
- How can we improve spatial aware learning? co generating captions for each frame with language decoder?
- How can we improve spatial aware learning? solve issue of Contrastive learning and position
- Extend evaluation []
- Add other datasets mentioned above []
- Ablation for size vision encoder, language encoder, temporal head