admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-24 10:16:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

increase stride

Browse Source
This commit is contained in:
Pepijn
2025-09-01 13:53:43 +02:00
parent a305f5f46a
commit f79fdf7205
2 changed files with 12 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/rlearn/configuration_rlearn.py
+2 -2
View File
@@ -46,7 +46,7 @@ class RLearNConfig(PreTrainedConfig):
# Sequence length, amount of past frames including current one to use in the temporal model # Sequence length, amount of past frames including current one to use in the temporal model
max_seq_len: int = 16 max_seq_len: int = 16
# Temporal sampling stride (2 = skip every other frame for wider temporal coverage) # 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 # Model dimensions and transformer
dim_model: int = 512 dim_model: int = 512
@@ -65,7 +65,7 @@ class RLearNConfig(PreTrainedConfig):
hl_gauss_num_bins: int = 25 # histogram resolution hl_gauss_num_bins: int = 25 # histogram resolution
# Inference-time subsampling and regularization # Inference-time subsampling and regularization
inference_stride: int = 2 inference_stride: int = 1 # in forward
frame_dropout_p: float = 0.10 frame_dropout_p: float = 0.10
# Training # Training
src/lerobot/policies/rlearn/modeling_rlearn.py
+10 -11
View File
@@ -118,16 +118,11 @@ class RLearNPolicy(PreTrainedPolicy):
self.reward_head = nn.Linear(config.dim_model, int(config.num_reward_bins)) self.reward_head = nn.Linear(config.dim_model, int(config.num_reward_bins))
self.hl_gauss_layer = None self.hl_gauss_layer = None
else: else:
# produce embeddings for HL-Gauss (or regression) # HL-Gauss expects per-bin logits; head outputs histogram-bin logits
self.reward_head = nn.Sequential( self.reward_head = nn.Linear(config.dim_model, int(config.hl_gauss_num_bins))
nn.Linear(config.dim_model, config.dim_model),
nn.GELU(),
nn.Dropout(config.dropout),
nn.Linear(config.dim_model, config.dim_model),
)
if HLGaussLayer is not None: if HLGaussLayer is not None:
self.hl_gauss_layer = HLGaussLayer( 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), use_regression=not bool(config.use_hl_gauss_loss),
hl_gauss_loss=dict( hl_gauss_loss=dict(
min_value=float(config.reward_min_value), min_value=float(config.reward_min_value),
@@ -380,7 +375,7 @@ class RLearNPolicy(PreTrainedPolicy):
predicted_rewards = torch.softmax(video_frame_logits, dim=-1) predicted_rewards = torch.softmax(video_frame_logits, dim=-1)
else: else:
# embeddings for HL-Gauss (or regression) # 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 # derive a scalar proxy for regularizers
raw_like_logits = torch.tanh(video_frame_embeds).mean(dim=-1) raw_like_logits = torch.tanh(video_frame_embeds).mean(dim=-1)
# predicted_rewards will be set after loss branch below # predicted_rewards will be set after loss branch below
@@ -441,9 +436,13 @@ class RLearNPolicy(PreTrainedPolicy):
else: else:
# HL-Gauss or regression # HL-Gauss or regression
if (self.hl_gauss_layer is not None) and (not self.hl_gauss_use_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 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: 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) pred_values = self.hl_gauss_layer(video_frame_embeds) # (B,T)
if video_mask is not None: if video_mask is not None: