diff --git a/src/lerobot/policies/fastwam/configuration_fastwam.py b/src/lerobot/policies/fastwam/configuration_fastwam.py index a3ef4f602..5907769fa 100644 --- a/src/lerobot/policies/fastwam/configuration_fastwam.py +++ b/src/lerobot/policies/fastwam/configuration_fastwam.py @@ -188,6 +188,14 @@ class FastWAMConfig(PreTrainedConfig): action_video_freq_ratio: int = 4 image_size: tuple[int, int] = (224, 448) context_len: int = 128 + + # Relative actions: converts absolute actions to relative (action -= state) during + # preprocessing, and reverses it at postprocessing. Requires `proprio_dim` (OBS_STATE). + use_relative_actions: bool = False + # Joint names to keep absolute (not converted to relative). Empty list = all dims relative. + relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"]) + # Populated at runtime from dataset metadata by make_policy (used to build the exclude mask). + action_feature_names: list[str] | None = None model_id: str = WAN22_MODEL_ID tokenizer_model_id: str = WAN_T5_TOKENIZER_ID text_encoder_model_id: str = WAN22_DIFFUSERS_MODEL_ID @@ -279,6 +287,10 @@ class FastWAMConfig(PreTrainedConfig): finally: self.pretrained_path = pretrained_path + @property + def chunk_size(self) -> int: + return self.action_horizon + def get_optimizer_preset(self) -> AdamWConfig: return AdamWConfig(lr=self.optimizer_lr, weight_decay=self.optimizer_weight_decay) diff --git a/src/lerobot/policies/fastwam/modeling_fastwam.py b/src/lerobot/policies/fastwam/modeling_fastwam.py index 10671e717..0563242b5 100644 --- a/src/lerobot/policies/fastwam/modeling_fastwam.py +++ b/src/lerobot/policies/fastwam/modeling_fastwam.py @@ -183,7 +183,9 @@ class FastWAMPolicy(PreTrainedPolicy): sample["proprio"] = state.unsqueeze(1) if state.ndim == 2 else state return sample - def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Any]]: + def forward( + self, batch: dict[str, Tensor], reduction: str = "mean" + ) -> tuple[Tensor, dict[str, Any]]: """Compute FastWAM training loss for a LeRobot batch. Args: @@ -191,15 +193,19 @@ class FastWAMPolicy(PreTrainedPolicy): (`video`, `action`, `context`, `context_mask`) or LeRobot keys that can be adapted (`observation.images.*`, `observation.state`, `action`, `action_is_pad`). + reduction (str): "mean" returns the scalar loss (default, backward + compatible); "none" returns per-sample losses of shape (batch_size,) + for sample weighting (RA-BC). Returns: - tuple[Tensor, dict[str, Any]]: The scalar loss to backprop, and a dict of - logging metrics (e.g. `loss_video`, `loss_action`) — the `(loss, output_dict)` - contract the LeRobot training loop expects. + tuple[Tensor, dict[str, Any]]: The loss to backprop (scalar for "mean", + per-sample (B,) for "none"), and a dict of logging metrics (e.g. + `loss_video`, `loss_action`) — the `(loss, output_dict)` contract the + LeRobot training loop expects. """ sample = self._batch_to_training_sample(batch) - loss, metrics = self.model.training_loss(sample) + loss, metrics = self.model.training_loss(sample, reduction=reduction) return loss, dict(metrics or {}) @torch.no_grad() diff --git a/src/lerobot/policies/fastwam/processor_fastwam.py b/src/lerobot/policies/fastwam/processor_fastwam.py index 31f3b9277..a8bf5c8c1 100644 --- a/src/lerobot/policies/fastwam/processor_fastwam.py +++ b/src/lerobot/policies/fastwam/processor_fastwam.py @@ -21,13 +21,16 @@ import torch from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.processor import ( + AbsoluteActionsProcessorStep, ActionProcessorStep, AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, + ProcessorStep, ProcessorStepRegistry, + RelativeActionsProcessorStep, RenameObservationsProcessorStep, UnnormalizerProcessorStep, policy_action_to_transition, @@ -105,10 +108,20 @@ def make_fastwam_pre_post_processors( # anyway) and unsafe across fine-tuning: its `resize_size` would be inherited from the base # checkpoint's camera geometry, not this dataset's, making the concatenation N_cameras x too wide. - input_steps = [ + # Shared relative-action step (OpenPI order: raw -> relative -> normalize -> model -> + # unnormalize -> absolute). The SAME instance is passed to AbsoluteActionsProcessorStep + # below so its cached raw state (set during preprocessing) flows to postprocessing. + relative_step = RelativeActionsProcessorStep( + enabled=config.use_relative_actions, + exclude_joints=getattr(config, "relative_exclude_joints", []), + action_names=getattr(config, "action_feature_names", None), + ) + + input_steps: list[ProcessorStep] = [ RenameObservationsProcessorStep(rename_map={}), AddBatchDimensionProcessorStep(), DeviceProcessorStep(device=config.device), + relative_step, NormalizerProcessorStep( features={**config.input_features, **config.output_features}, norm_map=config.normalization_mapping, @@ -116,12 +129,13 @@ def make_fastwam_pre_post_processors( device=config.device, ), ] - output_steps = [ + output_steps: list[ProcessorStep] = [ UnnormalizerProcessorStep( features=config.output_features, norm_map=config.normalization_mapping, stats=normalization_stats, ), + AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step), ] if config.toggle_action_dimensions: output_steps.append( diff --git a/src/lerobot/policies/fastwam/wan/modular.py b/src/lerobot/policies/fastwam/wan/modular.py index fac96776b..add9c82ef 100644 --- a/src/lerobot/policies/fastwam/wan/modular.py +++ b/src/lerobot/policies/fastwam/wan/modular.py @@ -1359,7 +1359,9 @@ class FastWAM(torch.nn.Module): pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre) return pred_video, pred_action - def _compute_training_video_loss(self, inputs, pred_video, target_video, timestep_video): + def _compute_training_video_loss( + self, inputs, pred_video, target_video, timestep_video, reduction: str = "mean" + ): include_initial_video_step = inputs["first_frame_latents"] is None if inputs["first_frame_latents"] is not None: pred_video = pred_video[:, :, 1:] @@ -1374,9 +1376,13 @@ class FastWAM(torch.nn.Module): loss_video_per_sample.device, dtype=loss_video_per_sample.dtype, ) - return (loss_video_per_sample * video_weight).mean() + weighted = loss_video_per_sample * video_weight + # reduction="none" returns the per-sample vector (B,) for sample weighting (RA-BC). + return weighted if reduction == "none" else weighted.mean() - def _compute_training_action_loss(self, inputs, pred_action, target_action, timestep_action): + def _compute_training_action_loss( + self, inputs, pred_action, target_action, timestep_action, reduction: str = "mean" + ): action_loss_token = functional.mse_loss( pred_action.float(), target_action.float(), reduction="none" ).mean(dim=2) @@ -1393,9 +1399,11 @@ class FastWAM(torch.nn.Module): action_loss_per_sample.device, dtype=action_loss_per_sample.dtype, ) - return (action_loss_per_sample * action_weight).mean() + weighted = action_loss_per_sample * action_weight + # reduction="none" returns the per-sample vector (B,) for sample weighting (RA-BC). + return weighted if reduction == "none" else weighted.mean() - def training_loss(self, sample, tiled: bool = False): + def training_loss(self, sample, tiled: bool = False, reduction: str = "mean"): inputs = self.build_inputs(sample, tiled=tiled) targets = self._sample_training_targets(inputs) pred_video, pred_action = self._run_training_mot(inputs=inputs, targets=targets) @@ -1404,17 +1412,20 @@ class FastWAM(torch.nn.Module): pred_video=pred_video, target_video=targets["target_video"], timestep_video=targets["timestep_video"], + reduction=reduction, ) loss_action = self._compute_training_action_loss( inputs=inputs, pred_action=pred_action, target_action=targets["target_action"], timestep_action=targets["timestep_action"], + reduction=reduction, ) + # With reduction="none" both terms are (B,), so loss_total is the per-sample loss (B,). loss_total = self.loss_lambda_video * loss_video + self.loss_lambda_action * loss_action loss_dict = { - "loss_video": self.loss_lambda_video * float(loss_video.detach().item()), - "loss_action": self.loss_lambda_action * float(loss_action.detach().item()), + "loss_video": self.loss_lambda_video * float(loss_video.detach().mean().item()), + "loss_action": self.loss_lambda_action * float(loss_action.detach().mean().item()), } return loss_total, loss_dict diff --git a/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py b/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py index 424ea7c63..bec865779 100644 --- a/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py +++ b/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py @@ -92,6 +92,15 @@ class LingBotVAConfig(PreTrainedConfig): # (un)normalization quantiles live in the checkpoint's ``policy_postprocessor.json``, not here. used_action_channel_ids: list[int] = field(default_factory=lambda: list(range(7))) + # Relative actions: converts absolute actions to relative (action -= state) during + # preprocessing, and reverses it at postprocessing. Requires the dataset to provide + # observation.state whose leading dims align 1:1 with the used action channels. + use_relative_actions: bool = False + # Joint names to keep absolute (not converted to relative). Empty list = all dims relative. + relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"]) + # Populated at runtime from dataset metadata by make_policy (used to build the exclude mask). + action_feature_names: list[str] | None = None + # Opt-in: VAE-decode predicted video latents to ``self.last_predicted_frames`` for saving MP4s. save_predicted_video: bool = False diff --git a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py index 0f70ad290..9953e9988 100644 --- a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py +++ b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py @@ -257,8 +257,12 @@ class LingBotVAPolicy(PreTrainedPolicy): "grid_id": grid_id, } - def _flow_matching_loss(self, input_dict, pred): - """Dual-stream flow-matching loss (port of upstream ``Trainer.compute_loss``).""" + def _flow_matching_loss(self, input_dict, pred, reduction: str = "mean"): + """Dual-stream flow-matching loss (port of upstream ``Trainer.compute_loss``). + + ``reduction="mean"`` returns scalar (latent_loss, action_loss); ``"none"`` returns + per-sample vectors of shape ``(B,)`` each (averaged over latent frames) for RA-BC. + """ latent_pred, action_pred = pred ld, ad = input_dict["latent_dict"], input_dict["action_dict"] action_pred = rearrange(action_pred, "b (f n) c -> b c f n 1", f=ad["targets"].shape[-3]) @@ -278,7 +282,8 @@ class LingBotVAPolicy(PreTrainedPolicy): latent_loss = ( (latent_loss * lw[:, None, :, None, None]).permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1) ) - latent_loss = (latent_loss.sum(dim=1) / (torch.ones_like(latent_loss).sum(dim=1) + 1e-6)).mean() + # per (batch*frame) mean over spatial/channel -> (B*F,) + latent_loss = latent_loss.sum(dim=1) / (torch.ones_like(latent_loss).sum(dim=1) + 1e-6) amask = ad["actions_mask"].float() action_loss = F.mse_loss(action_pred.float(), ad["targets"].float().detach(), reduction="none") @@ -286,10 +291,14 @@ class LingBotVAPolicy(PreTrainedPolicy): (action_loss * aw[:, None, :, None, None] * amask).permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1) ) amask_f = amask.permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1) - action_loss = (action_loss.sum(dim=1) / (amask_f.sum(dim=1) + 1e-6)).mean() - return latent_loss, action_loss + action_loss = action_loss.sum(dim=1) / (amask_f.sum(dim=1) + 1e-6) - def training_loss_from_streams(self, latents, actions, actions_mask, text_emb): + if reduction == "none": + # (B*F,) -> (B, F) -> (B,): per-sample losses for RA-BC weighting. + return latent_loss.reshape(bn, fn).mean(dim=1), action_loss.reshape(bn, fn).mean(dim=1) + return latent_loss.mean(), action_loss.mean() + + def training_loss_from_streams(self, latents, actions, actions_mask, text_emb, reduction: str = "mean"): """Core dual-stream training loss given prepared latents / actions / text embeddings. ``latents``: ``[B, in_channels, F, h, w]`` (normalized video latents). @@ -318,20 +327,24 @@ class LingBotVAPolicy(PreTrainedPolicy): "window_size": int(torch.randint(4, 65, (1,)).item()), } pred = self.transformer(input_dict, train_mode=True) - latent_loss, action_loss = self._flow_matching_loss(input_dict, pred) + latent_loss, action_loss = self._flow_matching_loss(input_dict, pred, reduction) + # reduction="none": latent_loss/action_loss are (B,) -> loss is per-sample (B,). loss = latent_loss + action_loss - return loss, {"latent_loss": latent_loss.item(), "action_loss": action_loss.item()} + return loss, {"latent_loss": latent_loss.mean().item(), "action_loss": action_loss.mean().item()} - def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]: + def forward(self, batch: dict[str, Tensor], reduction: str = "mean") -> tuple[Tensor, dict | None]: """Training forward: dual-stream flow-matching loss. Builds the (video-latent, action, text) training streams from a LeRobot batch (VAE-encoding the camera frames and UMT5-encoding the task), then runs the flow-matching dual-stream loss. Requires the policy to be built with ``attn_mode='flex'``. + + ``reduction="mean"`` returns the scalar loss (default); ``"none"`` returns per-sample + losses of shape ``(B,)`` for sample weighting (RA-BC). """ self._ensure_frozen_modules() latents, actions, actions_mask, text_emb = self._build_training_streams(batch) - return self.training_loss_from_streams(latents, actions, actions_mask, text_emb) + return self.training_loss_from_streams(latents, actions, actions_mask, text_emb, reduction=reduction) @torch.no_grad() def _build_training_streams(self, batch): diff --git a/src/lerobot/policies/lingbot_va/processor_lingbot_va.py b/src/lerobot/policies/lingbot_va/processor_lingbot_va.py index 119f77d5b..b081646bc 100644 --- a/src/lerobot/policies/lingbot_va/processor_lingbot_va.py +++ b/src/lerobot/policies/lingbot_va/processor_lingbot_va.py @@ -25,12 +25,14 @@ import torch from lerobot.configs.types import FeatureType, NormalizationMode from lerobot.processor import ( + AbsoluteActionsProcessorStep, AddBatchDimensionProcessorStep, DeviceProcessorStep, NormalizerProcessorStep, PolicyAction, PolicyProcessorPipeline, ProcessorStep, + RelativeActionsProcessorStep, RenameObservationsProcessorStep, UnnormalizerProcessorStep, ) @@ -52,9 +54,19 @@ def make_lingbot_va_pre_post_processors( ]: """Build the pre/post processor pipelines for LingBot-VA.""" + # Shared relative-action step (OpenPI order: raw -> relative -> normalize -> model -> + # unnormalize -> absolute). The SAME instance is passed to AbsoluteActionsProcessorStep + # below so its cached raw state (set during preprocessing) flows to postprocessing. + relative_step = RelativeActionsProcessorStep( + enabled=config.use_relative_actions, + exclude_joints=getattr(config, "relative_exclude_joints", []), + action_names=getattr(config, "action_feature_names", None), + ) + input_steps: list[ProcessorStep] = [ RenameObservationsProcessorStep(rename_map={}), AddBatchDimensionProcessorStep(), + relative_step, NormalizerProcessorStep( features={**config.input_features, **config.output_features}, norm_map=config.normalization_mapping, @@ -63,13 +75,16 @@ def make_lingbot_va_pre_post_processors( DeviceProcessorStep(device=config.device), ] - # Unnormalize actions from [-1, 1] to physical units (QUANTILES) using q01/q99 restored from the checkpoint. + # Unnormalize actions back to physical units. Config-driven norm_map (was hardcoded QUANTILES) + # so it stays symmetric with the preprocessor's NormalizerProcessorStep — required for + # use_relative_actions with ACTION=IDENTITY (and unchanged for QUANTILES runs). output_steps: list[ProcessorStep] = [ UnnormalizerProcessorStep( features=config.output_features, - norm_map={FeatureType.ACTION: NormalizationMode.QUANTILES}, + norm_map=config.normalization_mapping, stats=dataset_stats, ), + AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step), DeviceProcessorStep(device="cpu"), ] diff --git a/src/lerobot/processor/relative_action_processor.py b/src/lerobot/processor/relative_action_processor.py index 5b1039291..796c376f0 100644 --- a/src/lerobot/processor/relative_action_processor.py +++ b/src/lerobot/processor/relative_action_processor.py @@ -51,6 +51,12 @@ def to_relative_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) -> # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA. if state.device != actions.device or state.dtype != actions.dtype: state = state.to(device=actions.device, dtype=actions.dtype) + # When the observation is temporally stacked (e.g. LingBot-VA loads several obs steps via + # observation_delta_indices, giving state shape (B, T_obs, state_dim)), the relative reference + # is the CURRENT frame (delta == 0, i.e. index 0). Collapse to it so the offset broadcasts over + # the action horizon. pi0/pi05 pass a 2D (B, state_dim) state and are unaffected. + if state.ndim == 3: + state = state[:, 0] state_offset = state[..., :dims] * mask_t if actions.ndim == 3: state_offset = state_offset.unsqueeze(-2) @@ -73,6 +79,10 @@ def to_absolute_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) -> # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA. if state.device != actions.device or state.dtype != actions.dtype: state = state.to(device=actions.device, dtype=actions.dtype) + # Mirror to_relative_actions: collapse a temporally-stacked (B, T_obs, state_dim) state to the + # current frame (index 0) so the round-trip stays symmetric with the relative conversion. + if state.ndim == 3: + state = state[:, 0] state_offset = state[..., :dims] * mask_t if actions.ndim == 3: state_offset = state_offset.unsqueeze(-2)