From 43eedf62e488a6d9d5f7b180aa0fea90e8f289d2 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Sun, 31 Aug 2025 18:49:06 +0200 Subject: [PATCH] use dinov3 --- .../policies/rlearn/configuration_rlearn.py | 6 +- .../policies/rlearn/modeling_rlearn.py | 60 ++++++++++--------- 2 files changed, 36 insertions(+), 30 deletions(-) diff --git a/src/lerobot/policies/rlearn/configuration_rlearn.py b/src/lerobot/policies/rlearn/configuration_rlearn.py index d3e02d442..fc549a973 100644 --- a/src/lerobot/policies/rlearn/configuration_rlearn.py +++ b/src/lerobot/policies/rlearn/configuration_rlearn.py @@ -34,12 +34,12 @@ class RLearNConfig(PreTrainedConfig): Notes: - This follows the ReWiND paper architecture. It uses frozen vision/text encoders - (SigLIP2 for both vision and language) and trains a + (DINOv3 for vision, SigLIP2 for language) and trains a lightweight temporal aggregator + head. """ - # Encoders - Using SigLIP2 for both vision and text - vision_model_name: str = "google/siglip2-base-patch16-224" + # Encoders - Using DINOv3 for vision and SigLIP2 for text + vision_model_name: str = "facebook/dinov3-vitl16-pretrain-lvd1689m" text_model_name: str = "google/siglip2-base-patch16-224" freeze_backbones: bool = True diff --git a/src/lerobot/policies/rlearn/modeling_rlearn.py b/src/lerobot/policies/rlearn/modeling_rlearn.py index ae4b42e3c..2e5ad50ce 100644 --- a/src/lerobot/policies/rlearn/modeling_rlearn.py +++ b/src/lerobot/policies/rlearn/modeling_rlearn.py @@ -77,6 +77,7 @@ from __future__ import annotations import math import numpy as np +from contextlib import nullcontext from itertools import chain from operator import truediv @@ -118,13 +119,14 @@ class RLearNPolicy(PreTrainedPolicy): self.config = config self.episode_data_index = episode_data_index # Store episode boundaries for progress calculation - # Encoders - SigLIP2 for both vision and text - from transformers import AutoProcessor, AutoModel + # Encoders - DINOv3 for vision, SigLIP2 for text + from transformers import AutoProcessor, AutoModel, AutoImageProcessor - # Load SigLIP2 processors and models - self.vision_processor = AutoProcessor.from_pretrained(config.vision_model_name, use_fast=True) + # Load DINOv3 processor and model for vision + self.vision_processor = AutoImageProcessor.from_pretrained(config.vision_model_name) self.vision_model = AutoModel.from_pretrained(config.vision_model_name) + # Load SigLIP2 processor and model for text self.text_processor = AutoProcessor.from_pretrained(config.text_model_name, use_fast=True) self.text_model = AutoModel.from_pretrained(config.text_model_name) @@ -133,10 +135,11 @@ class RLearNPolicy(PreTrainedPolicy): self.vision_model = self.vision_model.to('cuda') self.text_model = self.text_model.to('cuda') - # Get hidden sizes from SigLIP2 config - vh = getattr(getattr(self.vision_model, 'config', None), 'vision_config', None) - self.vision_hidden = getattr(vh, 'hidden_size', 768) + # Get hidden sizes from models + # DINOv3-ViTL16 has hidden_size directly in config + self.vision_hidden = getattr(self.vision_model.config, 'hidden_size', 1024) # DINOv3-large default + # SigLIP2 text model th = getattr(getattr(self.text_model, 'config', None), 'text_config', None) self.text_hidden = getattr(th, 'hidden_size', 512) @@ -370,38 +373,41 @@ class RLearNPolicy(PreTrainedPolicy): # Optimized: Process tensor directly without numpy conversion device = next(self.vision_model.parameters()).device - # Normalize to [0, 1] if needed and ensure correct format for SigLIP2 + # Normalize to [0, 1] if needed and ensure correct format for DINOv3 if flat.dtype != torch.float32: flat = flat.float() if flat.max() > 1.0: flat = flat / 255.0 - # SigLIP2 expects images in [0, 1] range, RGB format - # Resize and normalize in batch - much faster than individual processing - try: - # Try direct tensor processing (faster path) - processed = self.vision_processor(images=flat, return_tensors="pt") - pixel_values = processed["pixel_values"].to(device) - except: - # Fallback to individual processing if needed, but optimized - # Convert entire batch to numpy at once (much faster) - flat_numpy = flat.permute(0, 2, 3, 1).cpu().numpy() # (BT, H, W, C) - images_list = [flat_numpy[i] for i in range(B * T)] - - processed = self.vision_processor(images=images_list, return_tensors="pt") - pixel_values = processed["pixel_values"].to(device) + # DINOv3 expects images in [0, 1] range, RGB format + # Convert tensor to list of PIL-like arrays for processor + flat_numpy = flat.permute(0, 2, 3, 1).cpu().numpy() # (BT, H, W, C) + images_list = [flat_numpy[i] for i in range(B * T)] - # Process in batch through vision model - vision_outputs = self.vision_model.vision_model(pixel_values=pixel_values) - cls_tokens = vision_outputs.last_hidden_state[:, 0] + # Process through DINOv3 processor and model + inputs = self.vision_processor(images=images_list, return_tensors="pt") + inputs = {k: v.to(device) for k, v in inputs.items()} + + # Process in batch through DINOv3 model + # Use inference mode for better performance when possible + context_manager = torch.inference_mode() if not self.training else nullcontext() + with context_manager: + vision_outputs = self.vision_model(**inputs) + + # Use pooler_output from DINOv3 (better than CLS token) + if hasattr(vision_outputs, 'pooler_output') and vision_outputs.pooler_output is not None: + vision_features_flat = vision_outputs.pooler_output # (BT, D) + else: + # Fallback to last hidden state CLS token if pooler_output not available + vision_features_flat = vision_outputs.last_hidden_state[:, 0] # (BT, D) # Reshape to (B, T, D) for analysis - vision_features = rearrange(cls_tokens, '(b t) d -> b t d', b=B, t=T) + vision_features = rearrange(vision_features_flat, '(b t) d -> b t d', b=B, t=T) # DEBUG: Analyze vision feature variability if self.training and torch.rand(1).item() < 0.05: # 5% of training steps with torch.no_grad(): - print(f"\nšŸ” VISION FEATURE DEBUG (B={B}, T={T}):") + print(f"\nšŸ” DINOv3 VISION FEATURE DEBUG (B={B}, T={T}):") # Check feature statistics feature_mean = vision_features.mean().item()