relative action for lingbot and fastwam

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)
 

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

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(

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
 

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
 

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

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

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)