increase stride

src/lerobot/policies/rlearn/configuration_rlearn.py

@@ -46,7 +46,7 @@ class RLearNConfig(PreTrainedConfig):
     # Sequence length, amount of past frames including current one to use in the temporal model
     max_seq_len: int = 16
     # Temporal sampling stride (2 = skip every other frame for wider temporal coverage)
-    temporal_sampling_stride: int = 2
+    temporal_sampling_stride: int = 3 # Open x mostly has fps 10, and rewind has seq len 16, ours is 30fps so 30/10 = 3 stride lenght to have same timeframe!
 
     # Model dimensions and transformer
     dim_model: int = 512
@@ -65,7 +65,7 @@ class RLearNConfig(PreTrainedConfig):
     hl_gauss_num_bins: int = 25                # histogram resolution
 
     # Inference-time subsampling and regularization
-    inference_stride: int = 2
+    inference_stride: int = 1 # in forward
     frame_dropout_p: float = 0.10
 
     # Training

src/lerobot/policies/rlearn/modeling_rlearn.py

@@ -118,16 +118,11 @@ class RLearNPolicy(PreTrainedPolicy):
             self.reward_head = nn.Linear(config.dim_model, int(config.num_reward_bins))
             self.hl_gauss_layer = None
         else:
-            # produce embeddings for HL-Gauss (or regression)
-            self.reward_head = nn.Sequential(
-                nn.Linear(config.dim_model, config.dim_model),
-                nn.GELU(),
-                nn.Dropout(config.dropout),
-                nn.Linear(config.dim_model, config.dim_model),
-            )
+            # HL-Gauss expects per-bin logits; head outputs histogram-bin logits
+            self.reward_head = nn.Linear(config.dim_model, int(config.hl_gauss_num_bins))
             if HLGaussLayer is not None:
                 self.hl_gauss_layer = HLGaussLayer(
-                    dim=config.dim_model,
+                    dim=int(config.hl_gauss_num_bins),
                     use_regression=not bool(config.use_hl_gauss_loss),
                     hl_gauss_loss=dict(
                         min_value=float(config.reward_min_value),
@@ -380,7 +375,7 @@ class RLearNPolicy(PreTrainedPolicy):
             predicted_rewards = torch.softmax(video_frame_logits, dim=-1)
         else:
             # embeddings for HL-Gauss (or regression)
-            video_frame_embeds = self.reward_head(frame_tokens)  # (B,T,D)
+            video_frame_embeds = self.reward_head(frame_tokens)  # (B,T,Bins)
             # derive a scalar proxy for regularizers
             raw_like_logits = torch.tanh(video_frame_embeds).mean(dim=-1)
             # predicted_rewards will be set after loss branch below
@@ -441,9 +436,13 @@ class RLearNPolicy(PreTrainedPolicy):
             else:
                 # HL-Gauss or regression
                 if (self.hl_gauss_layer is not None) and (not self.hl_gauss_use_regression):
-                    loss = self.hl_gauss_layer(video_frame_embeds, target, mask=video_mask)
+                    # Ensure targets within configured range
+                    t_min = float(self.config.reward_min_value)
+                    t_max = float(self.config.reward_max_value)
+                    target_clamped = target.clamp(t_min, t_max)
+                    loss = self.hl_gauss_layer(video_frame_embeds, target_clamped, mask=video_mask)
                     total_loss = loss
-                    predicted_rewards = self.hl_gauss_layer(video_frame_embeds).detach()
+                    predicted_rewards = self.hl_gauss_layer(video_frame_embeds)
                 elif (self.hl_gauss_layer is not None) and self.hl_gauss_use_regression:
                     pred_values = self.hl_gauss_layer(video_frame_embeds)  # (B,T)
                     if video_mask is not None: