From ffc3f9811a411b052e8b7461646b3495206312be Mon Sep 17 00:00:00 2001
From: Maxime Ellerbach <maxime.ellerbach@huggingface.co>
Date: Tue, 21 Jul 2026 08:31:00 +0000
Subject: [PATCH] vla-jepa relative actions

---
 .../vla_jepa/configuration_vla_jepa.py         |  8 ++++++++
 .../policies/vla_jepa/processor_vla_jepa.py    | 18 ++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py b/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py
index 8a30ee374..25f608c35 100644
--- a/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py
+++ b/src/lerobot/policies/vla_jepa/configuration_vla_jepa.py
@@ -56,6 +56,14 @@ class VLAJEPAConfig(PreTrainedConfig):
     action_dim: int = 7
     state_dim: int = 8
 
+    # Relative actions: converts absolute actions to relative (action -= state) during
+    # preprocessing, and reverses it at postprocessing. Requires `state_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
+
     num_action_tokens_per_timestep: int = 8
     num_embodied_action_tokens_per_instruction: int = 32
     num_inference_timesteps: int = 4
diff --git a/src/lerobot/policies/vla_jepa/processor_vla_jepa.py b/src/lerobot/policies/vla_jepa/processor_vla_jepa.py
index aed2b754b..e9cca0527 100644
--- a/src/lerobot/policies/vla_jepa/processor_vla_jepa.py
+++ b/src/lerobot/policies/vla_jepa/processor_vla_jepa.py
@@ -20,11 +20,13 @@ import torch
 
 from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig
 from lerobot.processor import (
+    AbsoluteActionsProcessorStep,
     EnvTransition,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
     ProcessorStepRegistry,
+    RelativeActionsProcessorStep,
     TransitionKey,
     UnnormalizerProcessorStep,
     make_default_policy_processor_steps,
@@ -109,10 +111,21 @@ def make_vla_jepa_pre_post_processors(
 ]:
     features = {**config.input_features, **config.output_features}
     steps = make_default_policy_processor_steps(config, dataset_stats)
+
+    # 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 = [
         steps.rename_observations,
         steps.add_batch_dim,
         steps.to_device,
+        relative_step,
         steps.normalize,
     ]
     output_steps: list[ProcessorStep] = []
@@ -131,6 +144,11 @@ def make_vla_jepa_pre_post_processors(
             stats=dataset_stats,
         )
     )
+    # Reverse the relative conversion on the unnormalized action, before gripper binarization.
+    # gripper is kept absolute by relative_exclude_joints, so the two steps touch disjoint dims.
+    output_steps.append(
+        AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step)
+    )
     if config.binarize_gripper_action:
         output_steps.append(
             BinarizeGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)