diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index cfb550ab2..2f8989a11 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -33,6 +33,7 @@ from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig from lerobot.policies.pi0.configuration_pi0 import PI0Config from lerobot.policies.pi05.configuration_pi05 import PI05Config from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.policies.rewind.configuration_rewind import ReWiNDConfig from lerobot.policies.sac.configuration_sac import SACConfig from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig @@ -293,6 +294,14 @@ def make_pre_post_processors( dataset_stats=kwargs.get("dataset_stats"), ) + elif isinstance(policy_cfg, ReWiNDConfig): + from lerobot.policies.rewind.processor_rewind import make_rewind_pre_post_processors + + processors = make_rewind_pre_post_processors( + config=policy_cfg, + dataset_stats=kwargs.get("dataset_stats"), + ) + else: raise NotImplementedError(f"Processor for policy type '{policy_cfg.type}' is not implemented.") diff --git a/src/lerobot/policies/rewind/__init__.py b/src/lerobot/policies/rewind/__init__.py index 22daf2d3e..a1aeff12a 100644 --- a/src/lerobot/policies/rewind/__init__.py +++ b/src/lerobot/policies/rewind/__init__.py @@ -18,19 +18,17 @@ from lerobot.policies.rewind.configuration_rewind import ReWiNDConfig from lerobot.policies.rewind.modeling_rewind import ( ReWiNDRewardModel, ReWiNDTransformer, - train_step_fn, - create_training_batch, - compute_progress_loss, - compute_misaligned_loss, +) +from lerobot.policies.rewind.processor_rewind import ( + ReWiNDEncodingProcessorStep, + make_rewind_pre_post_processors, ) __all__ = [ "ReWiNDConfig", "ReWiNDRewardModel", "ReWiNDTransformer", - "train_step_fn", - "create_training_batch", - "compute_progress_loss", - "compute_misaligned_loss", + "ReWiNDEncodingProcessorStep", + "make_rewind_pre_post_processors", ] diff --git a/src/lerobot/policies/rewind/configuration_rewind.py b/src/lerobot/policies/rewind/configuration_rewind.py index e55a2d0fd..a3fa0a035 100644 --- a/src/lerobot/policies/rewind/configuration_rewind.py +++ b/src/lerobot/policies/rewind/configuration_rewind.py @@ -14,9 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dataclasses import dataclass +from dataclasses import dataclass, field from lerobot.configs.policies import PreTrainedConfig +from lerobot.optim import OptimizerConfig +from lerobot.optim.schedulers import LRSchedulerConfig @PreTrainedConfig.register_subclass("rewind") @@ -54,6 +56,20 @@ class ReWiNDConfig(PreTrainedConfig): # Dropout dropout: float = 0.1 # Dropout rate for transformer + # Processor settings (for automatic preprocessing) + image_key: str = "observation.images.top" # Key for images in dataset + task_description: str = "perform the task" # Default task description + encode_on_the_fly: bool = True # Encode images/text during training + + # Features (required by PreTrainedPolicy) + input_features: dict = field(default_factory=lambda: { + "video_features": {"shape": [768], "dtype": "float32"}, + "text_features": {"shape": [384], "dtype": "float32"} + }) + output_features: dict = field(default_factory=lambda: { + "progress": {"shape": [1], "dtype": "float32"} + }) + def __post_init__(self): super().__post_init__() @@ -68,4 +84,24 @@ class ReWiNDConfig(PreTrainedConfig): if self.dropout < 0 or self.dropout >= 1: raise ValueError(f"dropout must be in [0, 1), got {self.dropout}") + + def get_optimizer_preset(self) -> OptimizerConfig: + """Get default optimizer configuration for ReWiND training.""" + return OptimizerConfig( + name="adamw", + lr=3e-4, + weight_decay=1e-4, + betas=(0.9, 0.999), + eps=1e-8, + grad_clip_norm=1.0 + ) + + def get_scheduler_preset(self) -> LRSchedulerConfig: + """Get default learning rate scheduler configuration.""" + return LRSchedulerConfig( + name="cosine", + warmup_steps=1000, + T_max=100000, # Will be overridden by training steps + eta_min=3e-5 + ) diff --git a/src/lerobot/policies/rewind/modeling_rewind.py b/src/lerobot/policies/rewind/modeling_rewind.py index 6b675f0bf..cb0d7be9f 100644 --- a/src/lerobot/policies/rewind/modeling_rewind.py +++ b/src/lerobot/policies/rewind/modeling_rewind.py @@ -27,6 +27,7 @@ from transformers import AutoModel, AutoTokenizer import torchvision.transforms as T from lerobot.policies.rewind.configuration_rewind import ReWiNDConfig +from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.datasets.video_sampler import sample_video_feature, sample_reverse_video_feature @@ -173,7 +174,7 @@ class ReWiNDTransformer(nn.Module): return progress_preds -class ReWiNDRewardModel(nn.Module): +class ReWiNDRewardModel(PreTrainedPolicy): """ ReWiND Reward Model for computing task completion rewards from video and text. @@ -183,9 +184,12 @@ class ReWiNDRewardModel(nn.Module): - ReWiNDTransformer for predicting task progress """ - def __init__(self, config: ReWiNDConfig): - super().__init__() + name = "rewind" + + def __init__(self, config: ReWiNDConfig, dataset_stats: dict | None = None): + super().__init__(config, dataset_stats) self.config = config + self.dataset_stats = dataset_stats self.device = torch.device(config.device if config.device else "cuda" if torch.cuda.is_available() else "cpu") # Initialize DINO encoder for images @@ -469,11 +473,119 @@ class ReWiNDRewardModel(nn.Module): def eval(self): """Set evaluation mode.""" return self.train(False) + + def parameters(self): + """Return trainable parameters (only ReWiND transformer, not encoders).""" + return self.rewind_transformer.parameters() + + def select_action(self, batch: dict[str, Tensor]) -> Tensor: + """ + This method is required by PreTrainedPolicy but not used for rewind. + The rewind model is not an actor and does not select actions. + """ + raise NotImplementedError("Rewind model does not select actions") + + def forward(self, batch): + """ + Forward pass compatible with lerobot training pipeline. + + Args: + batch: Dictionary containing: + - 'video_features': Pre-encoded video features (B, T, 768) + - 'text_features': Pre-encoded text features (B, 384) + - Optional: 'misaligned_video_features', 'misaligned_text_features' + + Returns: + loss: Total training loss + output_dict: Dictionary of loss components for logging + """ + # Use train_step_fn but without optimizer step (that's handled by training pipeline) + video_features = batch['video_features'].to(self.device) + text_features = batch['text_features'].to(self.device) + + batch_size = video_features.shape[0] + max_length = self.config.max_length + + # Process videos (with potential rewind augmentation) + import random + from lerobot.datasets.video_sampler import sample_video_feature, sample_reverse_video_feature + + processed_videos = [] + progress_targets = [] + + for i in range(batch_size): + if random.random() < 0.5: # 50% chance of rewind + # Apply video rewind augmentation + rewound_video, progress = sample_reverse_video_feature( + video_features[i], + max_length=max_length, + random_sample=True + ) + processed_videos.append(rewound_video) + progress_targets.append(progress) + else: + # Normal video sampling + sampled_video = sample_video_feature( + video_features[i], + max_length=max_length, + random_sample=True + ) + processed_videos.append(sampled_video) + # Linear progress from 0 to 1 + progress = torch.linspace(0, 1, max_length, device=self.device) + progress_targets.append(progress) + + processed_videos = torch.stack(processed_videos) + progress_targets = torch.stack(progress_targets) + + # Compute progress loss + progress_loss = compute_progress_loss( + self.rewind_transformer, + processed_videos, + text_features, + progress_targets + ) + + total_loss = progress_loss + output_dict = {'progress_loss': progress_loss.item()} + + # Compute misaligned loss if requested + if random.random() < 0.5: # 50% chance of adding misalignment loss + if 'misaligned_video_features' in batch and 'misaligned_text_features' in batch: + misaligned_videos = batch['misaligned_video_features'].to(self.device) + misaligned_texts = batch['misaligned_text_features'].to(self.device) + else: + # Create misaligned pairs by shuffling + shuffle_idx = torch.randperm(batch_size) + misaligned_videos = processed_videos[shuffle_idx] + misaligned_texts = text_features + + # Sample misaligned videos + misaligned_videos_sampled = [] + for i in range(batch_size): + sampled = sample_video_feature( + misaligned_videos[i], + max_length=max_length, + random_sample=True + ) + misaligned_videos_sampled.append(sampled) + misaligned_videos_sampled = torch.stack(misaligned_videos_sampled) + + misaligned_loss = compute_misaligned_loss( + self.rewind_transformer, + misaligned_videos_sampled, + misaligned_texts + ) + + total_loss = total_loss + misaligned_loss + output_dict['misaligned_loss'] = misaligned_loss.item() + + output_dict['total_loss'] = total_loss.item() + + return total_loss, output_dict -# Training utilities - - +# Loss utilities def compute_progress_loss( model: ReWiNDTransformer, video_features: torch.Tensor, @@ -538,171 +650,3 @@ def compute_misaligned_loss( loss = F.mse_loss(progress_preds, target_zeros) return loss - - -def train_step_fn( - model: ReWiNDRewardModel, - batch: Dict[str, torch.Tensor], - optimizer: torch.optim.Optimizer, - use_rewind: bool = True, - rewind_prob: float = 0.5, - misaligned_prob: float = 0.5, - gradient_clip: float = 1.0 -) -> Dict[str, float]: - """ - Perform a single training step for the ReWiND model. - - This function implements the training logic from the ReWiND paper, including: - - Progress prediction on aligned video-text pairs - - Video rewind augmentation for learning to decrease rewards - - Misaligned video-text pairs for learning to output zero rewards - - Args: - model: ReWiNDRewardModel instance - batch: Dictionary containing: - - 'video_features': Pre-computed video embeddings (batch_size, num_frames, 768) - - 'text_features': Pre-computed text embeddings (batch_size, 384) - - 'misaligned_video_features': Optional misaligned videos - - 'misaligned_text_features': Optional misaligned texts - optimizer: Optimizer for updating model parameters - use_rewind: Whether to use video rewind augmentation - rewind_prob: Probability of applying rewind to each sample - misaligned_prob: Probability of including misaligned loss - gradient_clip: Gradient clipping value - - Returns: - Dictionary of loss values for logging - """ - model.train() - optimizer.zero_grad() - - # Get features from batch - video_features = batch['video_features'].to(model.device) - text_features = batch['text_features'].to(model.device) - - batch_size = video_features.shape[0] - max_length = model.config.max_length - - # Process videos (with potential rewind augmentation) - processed_videos = [] - progress_targets = [] - - for i in range(batch_size): - if use_rewind and random.random() < rewind_prob: - # Apply video rewind augmentation - rewound_video, progress = sample_reverse_video_feature( - video_features[i], - max_length=max_length, - random_sample=True - ) - processed_videos.append(rewound_video) - progress_targets.append(progress) - else: - # Normal video sampling - sampled_video = sample_video_feature( - video_features[i], - max_length=max_length, - random_sample=True - ) - processed_videos.append(sampled_video) - # Linear progress from 0 to 1 - progress = torch.linspace(0, 1, max_length, device=model.device) - progress_targets.append(progress) - - processed_videos = torch.stack(processed_videos) - progress_targets = torch.stack(progress_targets) - - # Compute progress loss - progress_loss = compute_progress_loss( - model.rewind_transformer, - processed_videos, - text_features, - progress_targets - ) - - total_loss = progress_loss - losses = {'progress_loss': progress_loss.item()} - - # Compute misaligned loss if requested - if random.random() < misaligned_prob: - if 'misaligned_video_features' in batch and 'misaligned_text_features' in batch: - misaligned_videos = batch['misaligned_video_features'].to(model.device) - misaligned_texts = batch['misaligned_text_features'].to(model.device) - else: - # Create misaligned pairs by shuffling - shuffle_idx = torch.randperm(batch_size) - misaligned_videos = processed_videos[shuffle_idx] - misaligned_texts = text_features - - # Sample misaligned videos - misaligned_videos_sampled = [] - for i in range(batch_size): - sampled = sample_video_feature( - misaligned_videos[i], - max_length=max_length, - random_sample=True - ) - misaligned_videos_sampled.append(sampled) - misaligned_videos_sampled = torch.stack(misaligned_videos_sampled) - - misaligned_loss = compute_misaligned_loss( - model.rewind_transformer, - misaligned_videos_sampled, - misaligned_texts - ) - - total_loss = total_loss + misaligned_loss - losses['misaligned_loss'] = misaligned_loss.item() - - # Backward pass - total_loss.backward() - - # Gradient clipping - if gradient_clip > 0: - torch.nn.utils.clip_grad_norm_(model.rewind_transformer.parameters(), gradient_clip) - - # Optimizer step - optimizer.step() - - losses['total_loss'] = total_loss.item() - - return losses - - -def create_training_batch( - model: ReWiNDRewardModel, - videos: np.ndarray, - texts: List[str], - batch_size: int = 32, - encode_on_the_fly: bool = True -) -> Dict[str, torch.Tensor]: - """ - Create a training batch from raw videos and texts. - - Args: - model: ReWiNDRewardModel instance (for encoding if needed) - videos: Raw video frames (batch_size, num_frames, H, W, C) - texts: List of text descriptions - batch_size: Batch size for encoding - encode_on_the_fly: If True, encode videos and texts. If False, assume pre-encoded. - - Returns: - Dictionary containing video and text features - """ - if encode_on_the_fly: - # Encode videos using DINO - video_features = model.encode_images(videos) - video_features = torch.tensor(video_features, dtype=torch.float32) - - # Encode texts using MiniLM - text_features = model.encode_text(texts) - text_features = torch.tensor(text_features, dtype=torch.float32) - else: - # Assume videos and texts are already encoded - video_features = torch.tensor(videos, dtype=torch.float32) - text_features = torch.tensor(texts, dtype=torch.float32) - - return { - 'video_features': video_features, - 'text_features': text_features - } diff --git a/src/lerobot/policies/rewind/processor_rewind.py b/src/lerobot/policies/rewind/processor_rewind.py new file mode 100644 index 000000000..4837e7bad --- /dev/null +++ b/src/lerobot/policies/rewind/processor_rewind.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python + +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +from typing import Dict, Any, List, Optional +import numpy as np +import torch + +from lerobot.policies.rewind.configuration_rewind import ReWiNDConfig +from lerobot.policies.processor import ( + ProcessorStep, + PolicyProcessorPipeline, + PolicyAction, + DeviceProcessorStep, +) +from lerobot.policies.processor.transition import ( + policy_action_to_transition, + transition_to_policy_action, +) +from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME + +class ReWiNDEncodingProcessorStep(ProcessorStep): + """ + ProcessorStep that encodes images and text for ReWiND training. + + This step handles the DINO (image) and MiniLM (text) encoding that ReWiND needs. + """ + + def __init__( + self, + config: ReWiNDConfig, + image_key: str | None = None, + task_description: str | None = None, + ): + super().__init__() + self.config = config + self.image_key = image_key or config.image_key + self.task_description = task_description or config.task_description + + # Initialize encoders + self._init_encoders() + + def _init_encoders(self): + """Initialize DINO and MiniLM encoders.""" + from transformers import AutoModel, AutoTokenizer + + device = torch.device( + self.config.device if self.config.device + else "cuda" if torch.cuda.is_available() else "cpu" + ) + + logging.info("Initializing DINO encoder for ReWiND...") + self.dino_encoder = torch.hub.load("facebookresearch/dinov2", "dinov2_vitb14") + self.dino_encoder.to(device) + self.dino_encoder.eval() + + logging.info("Initializing MiniLM encoder for ReWiND...") + self.minilm_tokenizer = AutoTokenizer.from_pretrained( + "sentence-transformers/all-MiniLM-L12-v2" + ) + self.minilm_model = AutoModel.from_pretrained( + "sentence-transformers/all-MiniLM-L12-v2" + ) + self.minilm_model.to(device) + self.minilm_model.eval() + + self.device = device + + def __call__(self, batch: Dict[str, Any]) -> Dict[str, Any]: + """Encode images and text in the batch.""" + # Extract images + if self.image_key in batch: + images = batch[self.image_key] + + # Handle different image formats + if isinstance(images, torch.Tensor): + images = images.cpu().numpy() + + # Encode images + video_features = self._encode_images(images) + batch['video_features'] = video_features + + # Encode text + batch_size = len(batch.get('video_features', batch.get(list(batch.keys())[0]))) + task_descriptions = [self.task_description] * batch_size + text_features = self._encode_text(task_descriptions) + batch['text_features'] = text_features + + return batch + + @torch.no_grad() + def _encode_images(self, images: np.ndarray) -> torch.Tensor: + """Encode images using DINO.""" + from lerobot.policies.rewind.modeling_rewind import dino_load_image + + # Handle single frame case + if len(images.shape) == 4: + images = images[:, np.newaxis, ...] + single_frame = True + else: + single_frame = False + + batch_size, num_frames, C, H, W = images.shape + + # Convert to (B, T, H, W, C) + if C == 3: + images = images.transpose(0, 1, 3, 4, 2) + + all_embeddings = [] + + for video in images: + video_embeddings = [] + + # Convert to uint8 + if video.dtype != np.uint8: + video = (video * 255).astype(np.uint8) if video.max() <= 1.0 else video.astype(np.uint8) + + frames = [frame for frame in video] + episode_images_dino = [dino_load_image(frame) for frame in frames] + + # Batch process + for i in range(0, len(episode_images_dino), self.config.dino_batch_size): + dino_batch = torch.cat(episode_images_dino[i:i + self.config.dino_batch_size]) + dino_batch = dino_batch.to(self.device) + embeddings = self.dino_encoder(dino_batch).squeeze().detach().cpu() + + if embeddings.dim() == 1: + embeddings = embeddings.unsqueeze(0) + + video_embeddings.append(embeddings) + + video_embeddings = torch.cat(video_embeddings) + all_embeddings.append(video_embeddings) + + result = torch.stack(all_embeddings) + + if single_frame: + result = result.squeeze(1) + + return result + + @torch.no_grad() + def _encode_text(self, text: List[str]) -> torch.Tensor: + """Encode text using MiniLM.""" + from lerobot.policies.rewind.modeling_rewind import mean_pooling + + all_embeddings = [] + + for i in range(0, len(text), self.config.batch_size): + batch_text = text[i:i + self.config.batch_size] + + encoded_input = self.minilm_tokenizer( + batch_text, padding=True, truncation=True, return_tensors="pt" + ).to(self.device) + + model_output = self.minilm_model(**encoded_input) + text_embeddings = mean_pooling(model_output, encoded_input["attention_mask"]) + + all_embeddings.append(text_embeddings.cpu()) + + result = torch.cat(all_embeddings) + + return result + + +def make_rewind_pre_post_processors( + config: ReWiNDConfig, + dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, +) -> tuple[ + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], + PolicyProcessorPipeline[PolicyAction, PolicyAction], +]: + """ + Create pre-processor and post-processor pipelines for ReWiND. + + The pre-processing pipeline: + 1. Encodes images with DINO (768-dim) + 2. Encodes text with MiniLM (384-dim) + 3. Moves data to device + + The post-processing pipeline is minimal (just moves to CPU). + + Args: + config: ReWiND configuration + dataset_stats: Dataset statistics (not used for ReWiND) + + Returns: + Tuple of (preprocessor, postprocessor) pipelines + """ + input_steps = [ + ReWiNDEncodingProcessorStep(config=config), + DeviceProcessorStep(device=config.device), + ] + + output_steps = [ + DeviceProcessorStep(device="cpu"), + ] + + return ( + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]]( + steps=input_steps, + name=POLICY_PREPROCESSOR_DEFAULT_NAME, + ), + PolicyProcessorPipeline[PolicyAction, PolicyAction]( + steps=output_steps, + name=POLICY_POSTPROCESSOR_DEFAULT_NAME, + to_transition=policy_action_to_transition, + to_output=transition_to_policy_action, + ), + ) +