make rewind pretrained policy

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.")
 

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",
 ]
 

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
+        )
 

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
-    }

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,
+        ),
+    )
+