Add bundled OpenArm URDF and meshes for UMI evaluation

- Copy openarm_bimanual_pybullet.urdf and all referenced STL/DAE meshes
  into src/lerobot/robots/openarms/urdf/
- Update openarms_follower.py to prefer the bundled URDF over the
  external ~/Documents/openarm_description path
- Includes visual and collision meshes for arm, body, and end-effector
- Force-add .stl and .urdf files past .gitignore

examples/openarms/evaluate_relative.py

@@ -0,0 +1,321 @@
+#!/usr/bin/env python
+"""
+OpenArms Policy Evaluation with Relative Actions
+
+Two modes supported (based on training config):
+  Mode 1: Relative actions only (use_relative_state=False)
+    - Policy outputs relative action deltas
+    - State input is absolute
+  Mode 2: Relative actions + state (use_relative_state=True)
+    - Policy outputs relative action deltas  
+    - State input is also converted to relative
+
+Example usage:
+    python examples/openarms/evaluate_relative.py
+"""
+
+import time
+from pathlib import Path
+
+import torch
+
+from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.train import TrainPipelineConfig
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
+from lerobot.datasets.utils import build_dataset_frame, combine_feature_dicts
+from lerobot.policies.factory import make_policy, make_pre_post_processors
+from lerobot.policies.utils import predict_action
+from lerobot.processor import make_default_processors
+from lerobot.processor.core import RobotAction
+from lerobot.robots.openarms.config_openarms_follower import OpenArmsFollowerConfig
+from lerobot.robots.openarms.openarms_follower import OpenArmsFollower
+from lerobot.utils.constants import ACTION, OBS_STR
+from lerobot.utils.control_utils import init_keyboard_listener, precise_sleep
+from lerobot.utils.device_utils import get_safe_torch_device
+from lerobot.utils.relative_actions import (
+    convert_from_relative_actions_dict,
+    convert_state_to_relative,
+    PerTimestepNormalizer,
+)
+from lerobot.utils.utils import log_say
+from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
+
+
+# Configuration
+HF_MODEL_ID = "your-org/your-relative-policy"
+HF_EVAL_DATASET_ID = "your-org/your-eval-dataset"
+TASK_DESCRIPTION = "your task description"
+
+NUM_EPISODES = 1
+FPS = 30
+EPISODE_TIME_SEC = 300
+
+FOLLOWER_LEFT_PORT = "can0"
+FOLLOWER_RIGHT_PORT = "can1"
+
+CAMERA_CONFIG = {
+    "left_wrist": OpenCVCameraConfig(index_or_path="/dev/video5", width=640, height=480, fps=FPS),
+    "right_wrist": OpenCVCameraConfig(index_or_path="/dev/video1", width=640, height=480, fps=FPS),
+    "base": OpenCVCameraConfig(index_or_path="/dev/video3", width=640, height=480, fps=FPS),
+}
+
+
+def make_robot_action(action_values: dict, features: dict) -> RobotAction:
+    robot_action = {}
+    for key in features:
+        if key.startswith(ACTION + "."):
+            action_key = key.removeprefix(ACTION + ".")
+            if action_key in action_values:
+                robot_action[action_key] = action_values[action_key]
+    return robot_action
+
+
+def load_relative_config(model_path: Path | str) -> tuple[PerTimestepNormalizer | None, bool]:
+    """Load normalizer and relative_state setting from checkpoint."""
+    model_path = Path(model_path) if isinstance(model_path, str) else model_path
+    normalizer = None
+    use_relative_state = False
+    
+    # Try local path first
+    if model_path.exists():
+        stats_path = model_path / "relative_stats.pt"
+        if stats_path.exists():
+            normalizer = PerTimestepNormalizer.load(stats_path)
+            print(f"Loaded per-timestep stats from: {stats_path}")
+        
+        config_path = model_path / "train_config.json"
+        if config_path.exists():
+            cfg = TrainPipelineConfig.from_pretrained(model_path)
+            use_relative_state = getattr(cfg, "use_relative_state", False)
+    else:
+        # Try hub
+        try:
+            from huggingface_hub import hf_hub_download
+            stats_file = hf_hub_download(repo_id=str(model_path), filename="relative_stats.pt")
+            normalizer = PerTimestepNormalizer.load(stats_file)
+            print("Loaded per-timestep stats from hub")
+            
+            config_file = hf_hub_download(repo_id=str(model_path), filename="train_config.json")
+            cfg = TrainPipelineConfig.from_pretrained(Path(config_file).parent)
+            use_relative_state = getattr(cfg, "use_relative_state", False)
+        except Exception as e:
+            print(f"Warning: Could not load relative config: {e}")
+    
+    return normalizer, use_relative_state
+
+
+def inference_loop_relative(
+    robot,
+    policy,
+    preprocessor,
+    postprocessor,
+    dataset,
+    events,
+    fps: int,
+    control_time_s: float,
+    single_task: str,
+    display_data: bool = True,
+    state_key: str = "observation.state",
+    relative_normalizer: PerTimestepNormalizer | None = None,
+    use_relative_state: bool = False,
+):
+    """
+    Inference loop for relative action policies.
+    
+    If use_relative_state=True, also converts observation state to relative.
+    """
+    device = get_safe_torch_device(policy.config.device)
+    timestamp = 0
+    start_t = time.perf_counter()
+    
+    while timestamp < control_time_s:
+        loop_start = time.perf_counter()
+        
+        if events["exit_early"] or events["stop_recording"]:
+            break
+        
+        obs = robot.get_observation()
+        observation_frame = build_dataset_frame(dataset.features, obs, prefix=OBS_STR)
+        current_pos = {k: v for k, v in obs.items() if k.endswith(".pos")}
+        
+        # Convert state to relative if using full UMI mode
+        if use_relative_state and state_key in observation_frame:
+            state_tensor = observation_frame[state_key]
+            if isinstance(state_tensor, torch.Tensor):
+                observation_frame[state_key] = convert_state_to_relative(state_tensor)
+        
+        # Policy inference (outputs normalized relative actions)
+        action_values = predict_action(
+            observation=observation_frame,
+            policy=policy,
+            device=device,
+            preprocessor=preprocessor,
+            postprocessor=postprocessor,
+            use_amp=policy.config.use_amp,
+            task=single_task,
+            robot_type=robot.robot_type,
+        )
+        
+        # Unnormalize actions
+        if relative_normalizer is not None:
+            action_keys = [k for k in action_values.keys() if not k.startswith("task")]
+            action_tensor = torch.tensor([[action_values[k] for k in action_keys]])
+            action_tensor = action_tensor.unsqueeze(1)
+            action_unnorm = relative_normalizer.unnormalize(action_tensor)
+            for i, k in enumerate(action_keys):
+                action_values[k] = action_unnorm[0, 0, i].item()
+        
+        # Convert to absolute
+        relative_action = make_robot_action(action_values, dataset.features)
+        absolute_action = convert_from_relative_actions_dict(relative_action, current_pos)
+        
+        robot.send_action(absolute_action)
+        
+        if dataset is not None:
+            action_frame = build_dataset_frame(dataset.features, absolute_action, prefix=ACTION)
+            frame = {**observation_frame, **action_frame, "task": single_task}
+            dataset.add_frame(frame)
+        
+        if display_data:
+            log_rerun_data(observation=obs, action=absolute_action)
+        
+        dt = time.perf_counter() - loop_start
+        precise_sleep(1 / fps - dt)
+        timestamp = time.perf_counter() - start_t
+
+
+def main():
+    print("=" * 60)
+    print("  OpenArms Evaluation - Relative Actions")
+    print("=" * 60)
+    print(f"\nModel: {HF_MODEL_ID}")
+    print(f"Dataset: {HF_EVAL_DATASET_ID}")
+    print(f"Episodes: {NUM_EPISODES}, Duration: {EPISODE_TIME_SEC}s")
+    
+    # Load relative action config
+    relative_normalizer, use_relative_state = load_relative_config(HF_MODEL_ID)
+    mode = "actions + state" if use_relative_state else "actions only"
+    print(f"Mode: relative {mode}")
+    
+    # Setup robot
+    follower_config = OpenArmsFollowerConfig(
+        port_left=FOLLOWER_LEFT_PORT,
+        port_right=FOLLOWER_RIGHT_PORT,
+        can_interface="socketcan",
+        id="openarms_follower",
+        disable_torque_on_disconnect=True,
+        max_relative_target=10.0,
+        cameras=CAMERA_CONFIG,
+    )
+    
+    follower = OpenArmsFollower(follower_config)
+    follower.connect(calibrate=False)
+    
+    if not follower.is_connected:
+        raise RuntimeError("Robot failed to connect!")
+
+    teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
+    action_features_hw = {k: v for k, v in follower.action_features.items() if k.endswith(".pos")}
+    
+    dataset_features = combine_feature_dicts(
+        aggregate_pipeline_dataset_features(
+            pipeline=teleop_action_processor,
+            initial_features=create_initial_features(action=action_features_hw),
+            use_videos=True,
+        ),
+        aggregate_pipeline_dataset_features(
+            pipeline=robot_observation_processor,
+            initial_features=create_initial_features(observation=follower.observation_features),
+            use_videos=True,
+        ),
+    )
+    
+    dataset_path = Path.home() / ".cache" / "huggingface" / "lerobot" / HF_EVAL_DATASET_ID
+    if dataset_path.exists():
+        print(f"\nDataset exists at: {dataset_path}")
+        if input("Continue? (y/n): ").strip().lower() != 'y':
+            follower.disconnect()
+            return
+    
+    dataset = LeRobotDataset.create(
+        repo_id=HF_EVAL_DATASET_ID,
+        fps=FPS,
+        features=dataset_features,
+        robot_type=follower.name,
+        use_videos=True,
+        image_writer_processes=0,
+        image_writer_threads=12, 
+    )
+    
+    policy_config = PreTrainedConfig.from_pretrained(HF_MODEL_ID)
+    policy_config.pretrained_path = HF_MODEL_ID
+    policy = make_policy(policy_config, ds_meta=dataset.meta)
+    
+    preprocessor, postprocessor = make_pre_post_processors(
+        policy_cfg=policy.config,
+        pretrained_path=HF_MODEL_ID,
+        dataset_stats=dataset.meta.stats,
+        preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
+    )
+
+    listener, events = init_keyboard_listener()
+    init_rerun(session_name="openarms_eval_relative")
+    episode_idx = 0
+    
+    print("\nControls: ESC=stop, →=next episode, ←=rerecord")
+    
+    try:
+        while episode_idx < NUM_EPISODES and not events["stop_recording"]:
+            log_say(f"Episode {episode_idx + 1} of {NUM_EPISODES}")
+            
+            inference_loop_relative(
+                robot=follower,
+                policy=policy,
+                preprocessor=preprocessor,
+                postprocessor=postprocessor,
+                dataset=dataset,
+                events=events,
+                fps=FPS,
+                control_time_s=EPISODE_TIME_SEC,
+                single_task=TASK_DESCRIPTION,
+                display_data=True,
+                relative_normalizer=relative_normalizer,
+                use_relative_state=use_relative_state,
+            )
+            
+            if events.get("rerecord_episode", False):
+                log_say("Re-recording")
+                events["rerecord_episode"] = False
+                events["exit_early"] = False
+                dataset.clear_episode_buffer()
+                continue
+            
+            if dataset.episode_buffer is not None and dataset.episode_buffer.get("size", 0) > 0:
+                print(f"Saving episode {episode_idx + 1}...")
+                dataset.save_episode()
+                episode_idx += 1
+            
+            events["exit_early"] = False
+            
+            if not events["stop_recording"] and episode_idx < NUM_EPISODES:
+                input("Press ENTER for next episode...")
+        
+        print(f"\nDone! {episode_idx} episodes recorded")
+        log_say("Complete", blocking=True)
+    
+    except KeyboardInterrupt:
+        print("\n\nInterrupted")
+    
+    finally:
+        follower.disconnect()
+        if listener is not None:
+            listener.stop()
+        dataset.finalize()
+        print("Uploading to Hub...")
+        dataset.push_to_hub(private=True)
+
+
+if __name__ == "__main__":
+    main()

examples/openarms/evaluate_with_rtc.py

@@ -33,11 +33,6 @@ Example usage:
     python examples/openarms/evaluate_with_rtc.py \
         --rtc.execution_horizon=12 \
         --rtc.max_guidance_weight=10.0
-
-    # With action interpolation (policy at 30Hz, robot at 50Hz)
-    python examples/openarms/evaluate_with_rtc.py \
-        --action_interpolation_enabled=true \
-        --control_hz=50
 """
 
 import logging
@@ -87,8 +82,6 @@ DEFAULT_FPS = 30
 DEFAULT_EPISODE_TIME_SEC = 300
 DEFAULT_RESET_TIME_SEC = 60
 
-DEFAULT_CONTROL_HZ = 50
-
 DEFAULT_FOLLOWER_LEFT_PORT = "can0"
 DEFAULT_FOLLOWER_RIGHT_PORT = "can1"
 
@@ -174,9 +167,6 @@ class OpenArmsRTCEvalConfig(HubMixin):
     record_dataset: bool = True
     push_to_hub: bool = True
 
-    action_interpolation_enabled: bool = False
-    control_hz: float = DEFAULT_CONTROL_HZ
-
     use_torch_compile: bool = False
     torch_compile_backend: str = "inductor"
     torch_compile_mode: str = "default"
@@ -319,11 +309,6 @@ def get_actions_thread(
 # ============================================================================
 
 
-def _interpolate_actions(prev_action: Tensor, next_action: Tensor, alpha: float) -> Tensor:
-    """Linear interpolation between two action tensors."""
-    return prev_action + alpha * (next_action - prev_action)
-
-
 def actor_thread(
     robot: RobotWrapper,
     robot_action_processor,
@@ -339,101 +324,49 @@ def actor_thread(
     """Thread function to execute actions on the robot."""
     try:
         logger.info("[ACTOR] Starting actor thread")
-        logger.info(f"[ACTOR] interpolation={cfg.action_interpolation_enabled}, control_hz={cfg.control_hz}")
 
         action_count = 0
+        action_interval = 1.0 / cfg.fps
         action_keys = [k for k in robot.action_features.keys() if k.endswith(".pos")]
 
-        if cfg.action_interpolation_enabled:
-            control_interval = 1.0 / cfg.control_hz
-            interp_steps = int(cfg.control_hz / cfg.fps)
-        else:
-            control_interval = 1.0 / cfg.fps
-            interp_steps = 1
-
-        prev_action: Tensor | None = None
-        current_action: Tensor | None = None
-        interp_step = 0
-        last_dataset_frame_time = 0.0
-
         while not shutdown_event.is_set():
             if not episode_active.is_set():
-                prev_action = None
-                current_action = None
-                interp_step = 0
                 time.sleep(0.01)
                 continue
 
             start_time = time.perf_counter()
+            action = action_queue.get()
 
-            if cfg.action_interpolation_enabled:
-                if interp_step == 0 or current_action is None:
-                    new_action = action_queue.get()
-                    if new_action is not None:
-                        prev_action = current_action if current_action is not None else new_action.cpu()
-                        current_action = new_action.cpu()
-                        interp_step = 0
+            if action is not None:
+                action = action.cpu()
 
-                if current_action is not None:
-                    if prev_action is not None and interp_steps > 1:
-                        alpha = (interp_step + 1) / interp_steps
-                        action_to_send = _interpolate_actions(prev_action, current_action, alpha)
-                    else:
-                        action_to_send = current_action
+                action_dict = {}
+                for i, key in enumerate(action_keys):
+                    if i < len(action):
+                        action_dict[key] = action[i].item()
 
-                    action_dict = {}
-                    for i, key in enumerate(action_keys):
-                        if i < len(action_to_send):
-                            action_dict[key] = action_to_send[i].item()
+                action_processed = robot_action_processor((action_dict, None))
+                robot.send_action(action_processed)
 
-                    action_processed = robot_action_processor((action_dict, None))
-                    robot.send_action(action_processed)
-                    action_count += 1
+                if cfg.record_dataset and dataset is not None:
+                    with dataset_lock:
+                        obs = robot.get_observation()
+                        obs_processed = robot_observation_processor(obs)
+                        action_for_dataset = teleop_action_processor((action_dict, None))
 
-                    interp_step = (interp_step + 1) % interp_steps
+                        frame = {}
+                        for key, value in obs_processed.items():
+                            frame[f"observation.{key}"] = value
+                        for key, value in action_for_dataset.items():
+                            frame[f"action.{key}"] = value
+                        frame["task"] = cfg.task
 
-                    if cfg.record_dataset and dataset is not None:
-                        if time.perf_counter() - last_dataset_frame_time >= (1.0 / cfg.fps):
-                            last_dataset_frame_time = time.perf_counter()
-                            with dataset_lock:
-                                obs = robot.get_observation()
-                                obs_processed = robot_observation_processor(obs)
-                                action_for_dataset = teleop_action_processor((action_dict, None))
-                                frame = {}
-                                for key, value in obs_processed.items():
-                                    frame[f"observation.{key}"] = value
-                                for key, value in action_for_dataset.items():
-                                    frame[f"action.{key}"] = value
-                                frame["task"] = cfg.task
-                                dataset.add_frame(frame)
-            else:
-                action = action_queue.get()
-                if action is not None:
-                    action = action.cpu()
-                    action_dict = {}
-                    for i, key in enumerate(action_keys):
-                        if i < len(action):
-                            action_dict[key] = action[i].item()
+                        dataset.add_frame(frame)
 
-                    action_processed = robot_action_processor((action_dict, None))
-                    robot.send_action(action_processed)
-                    action_count += 1
-
-                    if cfg.record_dataset and dataset is not None:
-                        with dataset_lock:
-                            obs = robot.get_observation()
-                            obs_processed = robot_observation_processor(obs)
-                            action_for_dataset = teleop_action_processor((action_dict, None))
-                            frame = {}
-                            for key, value in obs_processed.items():
-                                frame[f"observation.{key}"] = value
-                            for key, value in action_for_dataset.items():
-                                frame[f"action.{key}"] = value
-                            frame["task"] = cfg.task
-                            dataset.add_frame(frame)
+                action_count += 1
 
             dt_s = time.perf_counter() - start_time
-            sleep_time = max(0, control_interval - dt_s - 0.001)
+            sleep_time = max(0, action_interval - dt_s - 0.001)
             if sleep_time > 0:
                 time.sleep(sleep_time)
 
@@ -501,9 +434,6 @@ def main(cfg: OpenArmsRTCEvalConfig):
     print(f"RTC Enabled: {cfg.rtc.enabled}")
     print(f"RTC Execution Horizon: {cfg.rtc.execution_horizon}")
     print(f"RTC Max Guidance Weight: {cfg.rtc.max_guidance_weight}")
-    print(f"Action Interpolation: {cfg.action_interpolation_enabled}")
-    if cfg.action_interpolation_enabled:
-        print(f"Control Hz: {cfg.control_hz}")
     print(f"Device: {cfg.device}")
     print("=" * 60)
 

examples/openarms/unify_task.py

@@ -1,152 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Unify all tasks in a dataset to a single task (modifies in-place).
-
-This script:
-1. Loads a dataset
-2. Sets all task_index to 0 and task description to "fold"
-3. Updates tasks.parquet and task_index in data files (in-place, no copying)
-
-Usage:
-    python examples/openarms/unify_task.py --repo-id lerobot-data-collection/level1_rac1
-"""
-
-from __future__ import annotations
-
-import argparse
-import logging
-from pathlib import Path
-
-import pandas as pd
-from tqdm import tqdm
-
-from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
-from lerobot.datasets.utils import (
-    DATA_DIR,
-    write_info,
-    write_tasks,
-)
-from lerobot.utils.constants import HF_LEROBOT_HOME
-
-
-# Single unified task
-UNIFIED_TASK = "fold"
-
-
-def unify_dataset_tasks(
-    repo_id: str,
-    root: Path | None = None,
-    push_to_hub: bool = False,
-) -> None:
-    """Unify all tasks in a dataset to a single task (modifies in-place).
-
-    Args:
-        repo_id: Dataset repository ID.
-        root: Optional root path for dataset.
-        push_to_hub: Whether to push the result to HuggingFace Hub.
-    """
-    input_root = root if root else HF_LEROBOT_HOME / repo_id
-    input_repo_id = repo_id
-
-    logging.info(f"Loading metadata from {repo_id}")
-
-    # Load source metadata
-    src_meta = LeRobotDatasetMetadata(repo_id, root=input_root)
-
-    logging.info(f"Source dataset: {src_meta.total_episodes} episodes, {src_meta.total_frames} frames")
-    logging.info(f"Original tasks: {len(src_meta.tasks)}")
-
-    # Modify in-place (input_root == output_root supported)
-    data_dir = input_root / DATA_DIR
-
-    # Process data files - set all task_index to 0
-    logging.info("Processing data files (in-place)...")
-    for parquet_file in tqdm(sorted(data_dir.rglob("*.parquet")), desc="Processing data"):
-        df = pd.read_parquet(parquet_file)
-        df["task_index"] = 0  # All tasks unified to index 0
-        df.to_parquet(parquet_file)
-
-    # Process episodes metadata - set all tasks to unified task
-    logging.info("Processing episodes metadata (in-place)...")
-    episodes_dir = input_root / "meta" / "episodes"
-    if episodes_dir.exists():
-        for parquet_file in tqdm(sorted(episodes_dir.rglob("*.parquet")), desc="Processing episodes"):
-            df = pd.read_parquet(parquet_file)
-            df["tasks"] = [[UNIFIED_TASK]] * len(df)  # All episodes get the unified task
-            df.to_parquet(parquet_file)
-    else:
-        logging.warning(f"No episodes directory found at {episodes_dir}, skipping")
-
-    # Update tasks.parquet with single task
-    logging.info(f"Creating single task: {UNIFIED_TASK}")
-    new_tasks = pd.DataFrame({"task_index": [0]}, index=[UNIFIED_TASK])
-    write_tasks(new_tasks, input_root)
-
-    # Update info.json
-    new_info = src_meta.info.copy()
-    new_info["total_tasks"] = 1
-    write_info(new_info, input_root)
-
-    logging.info(f"Dataset modified in-place at {input_root}")
-    logging.info(f"Task: {UNIFIED_TASK}")
-
-    if push_to_hub:
-        from lerobot.datasets.lerobot_dataset import LeRobotDataset
-
-        logging.info(f"Pushing {input_repo_id} to hub")
-        dataset = LeRobotDataset(input_repo_id, root=input_root)
-        dataset.push_to_hub(private=True)
-        logging.info("Push complete!")
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Unify all tasks in a dataset to a single task 'fold' (modifies in-place)."
-    )
-
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        required=True,
-        help="Dataset repository ID",
-    )
-    parser.add_argument(
-        "--root",
-        type=Path,
-        default=None,
-        help="Optional root path (defaults to HF_LEROBOT_HOME/repo_id)",
-    )
-    parser.add_argument(
-        "--push-to-hub",
-        action="store_true",
-        help="Push result to HuggingFace Hub",
-    )
-
-    args = parser.parse_args()
-
-    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
-
-    unify_dataset_tasks(
-        repo_id=args.repo_id,
-        root=args.root,
-        push_to_hub=args.push_to_hub,
-    )
-
-
-if __name__ == "__main__":
-    main()

scripts/unify_tasks.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python
+"""Unify all tasks in a dataset to a single task."""
+
+import argparse
+import json
+from pathlib import Path
+
+import pandas as pd
+
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import write_tasks
+
+
+def unify_tasks(repo_id: str, new_task: str):
+    """Set all episodes to use a single task."""
+    print(f"Loading dataset: {repo_id}")
+    dataset = LeRobotDataset(repo_id)
+    root = dataset.root
+    
+    print(f"Current tasks: {list(dataset.meta.tasks['task']) if dataset.meta.tasks is not None else []}")
+    
+    # 1. Update tasks.parquet to have only one task
+    tasks_df = pd.DataFrame({"task": [new_task]})
+    write_tasks(tasks_df, root)
+    print(f"Set single task: '{new_task}'")
+    
+    # 2. Update all data parquet files to set task_index=0
+    data_dir = root / "data"
+    parquet_files = sorted(data_dir.glob("*/*.parquet"))
+    for parquet_path in parquet_files:
+        df = pd.read_parquet(parquet_path)
+        df["task_index"] = 0
+        df.to_parquet(parquet_path)
+        print(f"Updated: {parquet_path.relative_to(root)}")
+    
+    # 3. Update info.json
+    info_path = root / "info.json"
+    with open(info_path) as f:
+        info = json.load(f)
+    info["total_tasks"] = 1
+    with open(info_path, "w") as f:
+        json.dump(info, f, indent=2)
+    
+    print(f"\nDone! All {dataset.meta.total_episodes} episodes now use task: '{new_task}'")
+    print(f"\nTo push: huggingface-cli upload {repo_id} {root} --repo-type dataset")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Unify all tasks in a dataset to a single task")
+    parser.add_argument("--repo_id", type=str, required=True, help="Dataset repo_id")
+    parser.add_argument("--task", type=str, required=True, help="New task description")
+    args = parser.parse_args()
+    
+    unify_tasks(args.repo_id, args.task)
+
+
+if __name__ == "__main__":
+    main()
+

src/lerobot/configs/train.py

@@ -66,6 +66,17 @@ class TrainPipelineConfig(HubMixin):
     eval: EvalConfig = field(default_factory=EvalConfig)
     wandb: WandBConfig = field(default_factory=WandBConfig)
 
+    # UMI-style relative actions with per-timestep normalization
+    # Mode 1: use_relative_actions=True, use_relative_state=False
+    #   - Actions: relative to current position + per-timestep normalized
+    #   - State: absolute (unchanged)
+    # Mode 2: use_relative_actions=True, use_relative_state=True (full UMI)
+    #   - Actions: relative to current position + per-timestep normalized  
+    #   - State: relative to current position (provides velocity info)
+    # Stats are computed automatically from first 1000 batches at training start
+    use_relative_actions: bool = False
+    use_relative_state: bool = False
+
     # RA-BC (Reward-Aligned Behavior Cloning) parameters
     use_rabc: bool = False  # Enable reward-weighted training
     rabc_progress_path: str | None = None  # Path to precomputed SARM progress parquet file

src/lerobot/robots/openarms/openarms_follower.py

@@ -96,16 +96,26 @@ class OpenArmsFollower(Robot):
         # Initialize Pinocchio robot model for dynamics (optional)
         self.pin_robot = None
         try:
-            # Load URDF - try external path first (with meshes), then repository
             import os
-            from os.path import expanduser, dirname
-            
-            # Try external URDF with meshes first
-            external_urdf_path = expanduser("~/Documents/openarm_description/openarm_bimanual_pybullet.urdf")
-            if os.path.exists(external_urdf_path):
+            from os.path import dirname
+
+            # Prefer the URDF bundled in the repository
+            repo_urdf_path = os.path.join(
+                dirname(__file__), "urdf", "openarm_bimanual_pybullet.urdf"
+            )
+            external_urdf_path = os.path.expanduser(
+                "~/Documents/openarm_description/openarm_bimanual_pybullet.urdf"
+            )
+
+            if os.path.exists(repo_urdf_path):
+                urdf_path = repo_urdf_path
+            elif os.path.exists(external_urdf_path):
                 urdf_path = external_urdf_path
+            else:
+                urdf_path = None
+
+            if urdf_path is not None:
                 urdf_dir = dirname(urdf_path)
-            
                 self.pin_robot = pin.RobotWrapper.BuildFromURDF(urdf_path, urdf_dir)
                 self.pin_robot.data = self.pin_robot.model.createData()
                 logger.info(f"Loaded OpenArms URDF for dynamics computation from {urdf_path}")

src/lerobot/robots/openarms/urdf/openarm_bimanual_pybullet.urdf

@@ -0,0 +1,618 @@
+<?xml version='1.0' encoding='utf-8'?>
+<robot name="openarm">
+  <link name="world" />
+  <joint name="openarm_body_world_joint" type="fixed">
+    <parent link="world" />
+    <child link="openarm_body_link0" />
+    <origin rpy="0 0 0" xyz="0 0 0" />
+  </joint>
+  <link name="openarm_body_link0">
+    <visual name="openarm_body_link0_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0" />
+      <geometry>
+        <mesh filename="./meshes/body/v10/visual/body_link0.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_body_link0_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0" />
+      <geometry>
+        <mesh filename="./meshes/body/v10/collision/body_link0_symp.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0" />
+      <mass value="13.89" />
+      <inertia ixx="1.653" ixy="0.0" ixz="0.0" iyy="1.653" iyz="0.0" izz="0.051" />
+    </inertial>
+  </link>
+  <joint name="openarm_left_openarm_body_link0_joint" type="fixed">
+    <parent link="openarm_body_link0" />
+    <child link="openarm_left_link0" />
+    <origin rpy="-1.5708 0 0" xyz="0.0 0.031 0.698" />
+  </joint>
+  <link name="openarm_left_link0">
+    <visual name="openarm_left_link0_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link0.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_left_link0_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link0_symp.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0009483362816297526 -0.0001580207020448382 0.03076860287587199" />
+      <mass value="1.1432284943239561" />
+      <inertia ixx="0.001128" ixy="-4e-06" ixz="-3.3e-05" iyy="0.000962" iyz="-7e-06" izz="0.00147" />
+    </inertial>
+  </link>
+  <link name="openarm_left_link1">
+    <visual name="openarm_left_link1_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0 0.0 -0.0625" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link1.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_left_link1_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0 0.0 -0.0625" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link1_symp.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0011467657911800769 -3.319987657026362e-05 0.05395284380736254" />
+      <mass value="1.1416684646202298" />
+      <inertia ixx="0.001567" ixy="-1e-06" ixz="-2.9e-05" iyy="0.001273" iyz="1e-06" izz="0.001016" />
+    </inertial>
+  </link>
+  <joint name="openarm_left_joint1" type="revolute">
+    <origin rpy="0 0 0" xyz="0.0 0.0 0.0625" />
+    <parent link="openarm_left_link0" />
+    <child link="openarm_left_link1" />
+    <axis xyz="0 0 1" />
+    <limit effort="40" lower="-3.490659" upper="1.3962629999999998" velocity="16.754666" />
+  </joint>
+  <link name="openarm_left_link2">
+    <visual name="openarm_left_link2_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0301 0.0 -0.1225" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link2.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_left_link2_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0301 0.0 -0.1225" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link2_symp.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="0.00839629182351943 2.0145102027597523e-08 0.03256649300522363" />
+      <mass value="0.2775092746011571" />
+      <inertia ixx="0.000359" ixy="1e-06" ixz="-0.000109" iyy="0.000376" iyz="1e-06" izz="0.000232" />
+    </inertial>
+  </link>
+  <joint name="openarm_left_joint2" type="revolute">
+    <origin rpy="-1.57079632679 0 0" xyz="-0.0301 0.0 0.06" />
+    <parent link="openarm_left_link1" />
+    <child link="openarm_left_link2" />
+    <axis xyz="-1 0 0" />
+    <limit effort="40" lower="-3.3161253267948965" upper="0.17453267320510335" velocity="16.754666" />
+  </joint>
+  <link name="openarm_left_link3">
+    <visual name="openarm_left_link3_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.18875" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link3.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_left_link3_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.18875" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link3_symp.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="-0.002104752099628911 -0.0005549085042607548 0.09047470545721961" />
+      <mass value="1.073863338202347" />
+      <inertia ixx="0.004372" ixy="1e-06" ixz="1.1e-05" iyy="0.004319" iyz="-3.6e-05" izz="0.000661" />
+    </inertial>
+  </link>
+  <joint name="openarm_left_joint3" type="revolute">
+    <origin rpy="0 0 0" xyz="0.0301 0.0 0.06625" />
+    <parent link="openarm_left_link2" />
+    <child link="openarm_left_link3" />
+    <axis xyz="0 0 1" />
+    <limit effort="27" lower="-1.570796" upper="1.570796" velocity="5.445426" />
+  </joint>
+  <link name="openarm_left_link4">
+    <visual name="openarm_left_link4_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0315 -0.3425" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link4.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_left_link4_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0315 -0.3425" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link4_symp.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0029006831074562967 -0.03030575826634669 0.06339637422196209" />
+      <mass value="0.6348534566833373" />
+      <inertia ixx="0.000623" ixy="-1e-06" ixz="-1.9e-05" iyy="0.000511" iyz="3.8e-05" izz="0.000334" />
+    </inertial>
+  </link>
+  <joint name="openarm_left_joint4" type="revolute">
+    <origin rpy="0 0 0" xyz="-0.0 0.0315 0.15375" />
+    <parent link="openarm_left_link3" />
+    <child link="openarm_left_link4" />
+    <axis xyz="0 1 0" />
+    <limit effort="27" lower="0.0" upper="2.443461" velocity="5.445426" />
+  </joint>
+  <link name="openarm_left_link5">
+    <visual name="openarm_left_link5_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.438" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link5.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_left_link5_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.438" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link5_symp.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="-0.003049665024221911 -0.0008866902457326625 0.043079803024980934" />
+      <mass value="0.6156588026168502" />
+      <inertia ixx="0.000423" ixy="-8e-06" ixz="6e-06" iyy="0.000445" iyz="-6e-06" izz="0.000324" />
+    </inertial>
+  </link>
+  <joint name="openarm_left_joint5" type="revolute">
+    <origin rpy="0 0 0" xyz="0.0 -0.0315 0.0955" />
+    <parent link="openarm_left_link4" />
+    <child link="openarm_left_link5" />
+    <axis xyz="0 0 1" />
+    <limit effort="7" lower="-1.570796" upper="1.570796" velocity="20.943946" />
+  </joint>
+  <link name="openarm_left_link6">
+    <visual name="openarm_left_link6_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0375 -0.0 -0.5585" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link6.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_left_link6_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0375 -0.0 -0.5585" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link6_symp.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="-0.037136587005447405 -0.00033230528343419053 -9.498374522309838e-05" />
+      <mass value="0.475202773187987" />
+      <inertia ixx="0.000143" ixy="1e-06" ixz="1e-06" iyy="0.000157" iyz="1e-06" izz="0.000159" />
+    </inertial>
+  </link>
+  <joint name="openarm_left_joint6" type="revolute">
+    <origin rpy="0 0 0" xyz="0.0375 0.0 0.1205" />
+    <parent link="openarm_left_link5" />
+    <child link="openarm_left_link6" />
+    <axis xyz="1 0 0" />
+    <limit effort="7" lower="-0.785398" upper="0.785398" velocity="20.943946" />
+  </joint>
+  <link name="openarm_left_link7">
+    <visual name="openarm_left_link7_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0 -0.5585" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link7.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_left_link7_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0 -0.5585" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link7_symp.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="6.875510271106056e-05 -0.01266175250761268 0.06951945409987448" />
+      <mass value="0.4659771327380578" />
+      <inertia ixx="0.000639" ixy="1e-06" ixz="1e-06" iyy="0.000497" iyz="8.9e-05" izz="0.000342" />
+    </inertial>
+  </link>
+  <joint name="openarm_left_joint7" type="revolute">
+    <origin rpy="0 0 0" xyz="-0.0375 0.0 0.0" />
+    <parent link="openarm_left_link6" />
+    <child link="openarm_left_link7" />
+    <axis xyz="0 -1 0" />
+    <limit effort="7" lower="-1.570796" upper="1.570796" velocity="20.943946" />
+  </joint>
+  <joint name="openarm_right_openarm_body_link0_joint" type="fixed">
+    <parent link="openarm_body_link0" />
+    <child link="openarm_right_link0" />
+    <origin rpy="1.5708 0 0" xyz="0.0 -0.031 0.698" />
+  </joint>
+  <link name="openarm_right_link0">
+    <visual name="openarm_right_link0_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link0.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_right_link0_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link0_symp.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0009483362816297526 0.0001580207020448382 0.03076860287587199" />
+      <mass value="1.1432284943239561" />
+      <inertia ixx="0.001128" ixy="-4e-06" ixz="-3.3e-05" iyy="0.000962" iyz="-7e-06" izz="0.00147" />
+    </inertial>
+  </link>
+  <link name="openarm_right_link1">
+    <visual name="openarm_right_link1_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0 0.0 -0.0625" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link1.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_right_link1_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0 0.0 -0.0625" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link1_symp.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0011467657911800769 3.319987657026362e-05 0.05395284380736254" />
+      <mass value="1.1416684646202298" />
+      <inertia ixx="0.001567" ixy="-1e-06" ixz="-2.9e-05" iyy="0.001273" iyz="1e-06" izz="0.001016" />
+    </inertial>
+  </link>
+  <joint name="openarm_right_joint1" type="revolute">
+    <origin rpy="0 0 0" xyz="0.0 0.0 0.0625" />
+    <parent link="openarm_right_link0" />
+    <child link="openarm_right_link1" />
+    <axis xyz="0 0 1" />
+    <limit effort="40" lower="-1.396263" upper="3.490659" velocity="16.754666" />
+  </joint>
+  <link name="openarm_right_link2">
+    <visual name="openarm_right_link2_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0301 0.0 -0.1225" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link2.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_right_link2_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0301 0.0 -0.1225" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link2_symp.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="0.00839629182351943 -2.0145102027597523e-08 0.03256649300522363" />
+      <mass value="0.2775092746011571" />
+      <inertia ixx="0.000359" ixy="1e-06" ixz="-0.000109" iyy="0.000376" iyz="1e-06" izz="0.000232" />
+    </inertial>
+  </link>
+  <joint name="openarm_right_joint2" type="revolute">
+    <origin rpy="1.57079632679 0 0" xyz="-0.0301 0.0 0.06" />
+    <parent link="openarm_right_link1" />
+    <child link="openarm_right_link2" />
+    <axis xyz="-1 0 0" />
+    <limit effort="40" lower="-0.17453267320510335" upper="3.3161253267948965" velocity="16.754666" />
+  </joint>
+  <link name="openarm_right_link3">
+    <visual name="openarm_right_link3_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.18875" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link3.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_right_link3_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.18875" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link3_symp.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="-0.002104752099628911 0.0005549085042607548 0.09047470545721961" />
+      <mass value="1.073863338202347" />
+      <inertia ixx="0.004372" ixy="1e-06" ixz="1.1e-05" iyy="0.004319" iyz="-3.6e-05" izz="0.000661" />
+    </inertial>
+  </link>
+  <joint name="openarm_right_joint3" type="revolute">
+    <origin rpy="0 0 0" xyz="0.0301 0.0 0.06625" />
+    <parent link="openarm_right_link2" />
+    <child link="openarm_right_link3" />
+    <axis xyz="0 0 1" />
+    <limit effort="27" lower="-1.570796" upper="1.570796" velocity="5.445426" />
+  </joint>
+  <link name="openarm_right_link4">
+    <visual name="openarm_right_link4_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0315 -0.3425" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link4.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_right_link4_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0315 -0.3425" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link4_symp.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0029006831074562967 -0.03030575826634669 0.06339637422196209" />
+      <mass value="0.6348534566833373" />
+      <inertia ixx="0.000623" ixy="-1e-06" ixz="-1.9e-05" iyy="0.000511" iyz="3.8e-05" izz="0.000334" />
+    </inertial>
+  </link>
+  <joint name="openarm_right_joint4" type="revolute">
+    <origin rpy="0 0 0" xyz="-0.0 0.0315 0.15375" />
+    <parent link="openarm_right_link3" />
+    <child link="openarm_right_link4" />
+    <axis xyz="0 1 0" />
+    <limit effort="27" lower="0.0" upper="2.443461" velocity="5.445426" />
+  </joint>
+  <link name="openarm_right_link5">
+    <visual name="openarm_right_link5_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.438" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link5.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_right_link5_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0 -0.0 -0.438" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link5_symp.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="-0.003049665024221911 0.0008866902457326625 0.043079803024980934" />
+      <mass value="0.6156588026168502" />
+      <inertia ixx="0.000423" ixy="-8e-06" ixz="6e-06" iyy="0.000445" iyz="-6e-06" izz="0.000324" />
+    </inertial>
+  </link>
+  <joint name="openarm_right_joint5" type="revolute">
+    <origin rpy="0 0 0" xyz="0.0 -0.0315 0.0955" />
+    <parent link="openarm_right_link4" />
+    <child link="openarm_right_link5" />
+    <axis xyz="0 0 1" />
+    <limit effort="7" lower="-1.570796" upper="1.570796" velocity="20.943946" />
+  </joint>
+  <link name="openarm_right_link6">
+    <visual name="openarm_right_link6_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0375 -0.0 -0.5585" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link6.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_right_link6_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="-0.0375 -0.0 -0.5585" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link6_symp.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="-0.037136587005447405 0.00033230528343419053 -9.498374522309838e-05" />
+      <mass value="0.475202773187987" />
+      <inertia ixx="0.000143" ixy="1e-06" ixz="1e-06" iyy="0.000157" iyz="1e-06" izz="0.000159" />
+    </inertial>
+  </link>
+  <joint name="openarm_right_joint6" type="revolute">
+    <origin rpy="0 0 0" xyz="0.0375 0.0 0.1205" />
+    <parent link="openarm_right_link5" />
+    <child link="openarm_right_link6" />
+    <axis xyz="1 0 0" />
+    <limit effort="7" lower="-0.785398" upper="0.785398" velocity="20.943946" />
+  </joint>
+  <link name="openarm_right_link7">
+    <visual name="openarm_right_link7_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0 -0.5585" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/visual/link7.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_right_link7_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 -0.0 -0.5585" />
+      <geometry>
+        <mesh filename="./meshes/arm/v10/collision/link7_symp.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0.0 0.0 0.0" xyz="6.875510271106056e-05 0.01266175250761268 0.06951945409987448" />
+      <mass value="0.4659771327380578" />
+      <inertia ixx="0.000639" ixy="1e-06" ixz="1e-06" iyy="0.000497" iyz="8.9e-05" izz="0.000342" />
+    </inertial>
+  </link>
+  <joint name="openarm_right_joint7" type="revolute">
+    <origin rpy="0 0 0" xyz="-0.0375 0.0 0.0" />
+    <parent link="openarm_right_link6" />
+    <child link="openarm_right_link7" />
+    <axis xyz="0 1 0" />
+    <limit effort="7" lower="-1.570796" upper="1.570796" velocity="20.943946" />
+  </joint>
+  <link name="openarm_left_hand">
+    <visual name="openarm_left_hand_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 -0.6585" />
+      <geometry>
+        <mesh filename="./meshes/ee/openarm_hand/visual/hand.dae" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_left_hand_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 -0.6585" />
+      <geometry>
+        <mesh filename="./meshes/ee/openarm_hand/collision/hand.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.0 0.002 0.03" />
+      <mass value="0.35" />
+      <inertia ixx="0.0002473" ixy="1e-06" ixz="1e-06" iyy="1.763e-05" iyz="1e-06" izz="0.0002521" />
+    </inertial>
+  </link>
+  <joint name="left_openarm_hand_joint" type="fixed">
+    <parent link="openarm_left_link7" />
+    <child link="openarm_left_hand" />
+    <origin rpy="0 0 0" xyz="0 -0.0 0.1001" />
+  </joint>
+  <link name="openarm_left_hand_tcp">
+    <inertial>
+      <origin xyz="0 0 0" rpy="0 0 0" />
+      <mass value="0.001" />
+      <inertia ixx="0.000001" ixy="0.0" ixz="0.0" iyy="0.000001" iyz="0.0" izz="0.000001" />
+    </inertial>
+  </link>
+  <joint name="openarm_left_hand_tcp_joint" type="fixed">
+    <origin rpy="0 0 0" xyz="0 -0.0 0.08" />
+    <parent link="openarm_left_hand" />
+    <child link="openarm_left_hand_tcp" />
+  </joint>
+  <link name="openarm_left_left_finger">
+    <visual name="openarm_left_left_finger_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 -0.05 -0.673001" />
+      <geometry>
+        <mesh filename="./meshes/ee/openarm_hand/visual/finger.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_left_left_finger_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 -0.05 -0.673001" />
+      <geometry>
+        <mesh filename="./meshes/ee/openarm_hand/collision/finger.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.0064528 0.01702 0.0219685" />
+      <mass value="0.03602545343277134" />
+      <inertia ixx="2.3749999999999997e-06" ixy="1e-06" ixz="1e-06" iyy="2.3749999999999997e-06" iyz="1e-06" izz="7.5e-07" />
+    </inertial>
+  </link>
+  <link name="openarm_left_right_finger">
+    <visual name="openarm_left_right_finger_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.05 -0.673001" />
+      <geometry>
+        <mesh filename="./meshes/ee/openarm_hand/visual/finger.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_left_right_finger_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.05 -0.673001" />
+      <geometry>
+        <mesh filename="./meshes/ee/openarm_hand/collision/finger.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.0064528 -0.01702 0.0219685" />
+      <mass value="0.03602545343277134" />
+      <inertia ixx="2.3749999999999997e-06" ixy="1e-06" ixz="1e-06" iyy="2.3749999999999997e-06" iyz="1e-06" izz="7.5e-07" />
+    </inertial>
+  </link>
+  <joint name="openarm_left_finger_joint1" type="prismatic">
+    <parent link="openarm_left_hand" />
+    <child link="openarm_left_right_finger" />
+    <origin rpy="0 0 0" xyz="0 -0.006 0.015" />
+    <axis xyz="0 -1 0" />
+    <limit effort="333" lower="0.0" upper="0.044" velocity="10.0" />
+  </joint>
+  <joint name="openarm_left_finger_joint2" type="prismatic">
+    <parent link="openarm_left_hand" />
+    <child link="openarm_left_left_finger" />
+    <origin rpy="0 0 0" xyz="0 0.006 0.015" />
+    <axis xyz="0 1 0" />
+    <limit effort="333" lower="0.0" upper="0.044" velocity="10.0" />
+    <mimic joint="openarm_left_finger_joint1" />
+  </joint>
+  <link name="openarm_right_hand">
+    <visual name="openarm_right_hand_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 -0.6585" />
+      <geometry>
+        <mesh filename="./meshes/ee/openarm_hand/visual/hand.dae" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_right_hand_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 -0.6585" />
+      <geometry>
+        <mesh filename="./meshes/ee/openarm_hand/collision/hand.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.0 0.002 0.03" />
+      <mass value="0.35" />
+      <inertia ixx="0.0002473" ixy="1e-06" ixz="1e-06" iyy="1.763e-05" iyz="1e-06" izz="0.0002521" />
+    </inertial>
+  </link>
+  <joint name="right_openarm_hand_joint" type="fixed">
+    <parent link="openarm_right_link7" />
+    <child link="openarm_right_hand" />
+    <origin rpy="0 0 0" xyz="0 -0.0 0.1001" />
+  </joint>
+  <link name="openarm_right_hand_tcp">
+    <inertial>
+      <origin xyz="0 0 0" rpy="0 0 0" />
+      <mass value="0.001" />
+      <inertia ixx="0.000001" ixy="0.0" ixz="0.0" iyy="0.000001" iyz="0.0" izz="0.000001" />
+    </inertial>
+  </link>
+  <joint name="openarm_right_hand_tcp_joint" type="fixed">
+    <origin rpy="0 0 0" xyz="0 -0.0 0.08" />
+    <parent link="openarm_right_hand" />
+    <child link="openarm_right_hand_tcp" />
+  </joint>
+  <link name="openarm_right_left_finger">
+    <visual name="openarm_right_left_finger_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 -0.05 -0.673001" />
+      <geometry>
+        <mesh filename="./meshes/ee/openarm_hand/visual/finger.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_right_left_finger_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 -0.05 -0.673001" />
+      <geometry>
+        <mesh filename="./meshes/ee/openarm_hand/collision/finger.stl" scale="0.001 0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.0064528 0.01702 0.0219685" />
+      <mass value="0.03602545343277134" />
+      <inertia ixx="2.3749999999999997e-06" ixy="1e-06" ixz="1e-06" iyy="2.3749999999999997e-06" iyz="1e-06" izz="7.5e-07" />
+    </inertial>
+  </link>
+  <link name="openarm_right_right_finger">
+    <visual name="openarm_right_right_finger_visual">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.05 -0.673001" />
+      <geometry>
+        <mesh filename="./meshes/ee/openarm_hand/visual/finger.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </visual>
+    <collision name="openarm_right_right_finger_collision">
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.05 -0.673001" />
+      <geometry>
+        <mesh filename="./meshes/ee/openarm_hand/collision/finger.stl" scale="0.001 -0.001 0.001" />
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.0064528 -0.01702 0.0219685" />
+      <mass value="0.03602545343277134" />
+      <inertia ixx="2.3749999999999997e-06" ixy="1e-06" ixz="1e-06" iyy="2.3749999999999997e-06" iyz="1e-06" izz="7.5e-07" />
+    </inertial>
+  </link>
+  <joint name="openarm_right_finger_joint1" type="prismatic">
+    <parent link="openarm_right_hand" />
+    <child link="openarm_right_right_finger" />
+    <origin rpy="0 0 0" xyz="0 -0.006 0.015" />
+    <axis xyz="0 -1 0" />
+    <limit effort="333" lower="0.0" upper="0.044" velocity="10.0" />
+  </joint>
+  <joint name="openarm_right_finger_joint2" type="prismatic">
+    <parent link="openarm_right_hand" />
+    <child link="openarm_right_left_finger" />
+    <origin rpy="0 0 0" xyz="0 0.006 0.015" />
+    <axis xyz="0 1 0" />
+    <limit effort="333" lower="0.0" upper="0.044" velocity="10.0" />
+    <mimic joint="openarm_right_finger_joint1" />
+  </joint>
+</robot>

src/lerobot/scripts/lerobot_train.py

@@ -13,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import gc
 import logging
 import time
 from contextlib import nullcontext
@@ -46,6 +47,10 @@ from lerobot.utils.train_utils import (
     save_checkpoint,
     update_last_checkpoint,
 )
+from lerobot.utils.relative_actions import (
+    convert_to_relative_actions,
+    compute_global_relative_stats,
+)
 from lerobot.utils.utils import (
     format_big_number,
     has_method,
@@ -298,6 +303,61 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
             device=device,
         )
 
+    # Compute relative action/state stats and hotswap them into the normalizer
+    raw_state_key = None
+    if cfg.use_relative_actions:
+        from lerobot.processor.normalize_processor import hotswap_stats
+
+        mode = "actions + state" if cfg.use_relative_state else "actions only"
+        cfg.output_dir.mkdir(parents=True, exist_ok=True)
+        stats_path = cfg.output_dir / "relative_stats.pt"
+
+        reverse_rename = {v: k for k, v in cfg.rename_map.items()} if cfg.rename_map else {}
+        raw_state_key = reverse_rename.get("observation.state", "observation.state")
+
+        if is_main_process:
+            logging.info(colored(f"Relative mode: {mode}", "cyan", attrs=["bold"]))
+
+            if stats_path.exists():
+                logging.info(f"Loading pre-computed relative stats from: {stats_path}")
+            else:
+                logging.info("Computing global relative stats (first 1000 batches)...")
+                stats_dataset = make_dataset(cfg)
+                temp_loader = torch.utils.data.DataLoader(
+                    stats_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=0
+                )
+                rel_stats = compute_global_relative_stats(
+                    temp_loader, state_key=raw_state_key,
+                    convert_state=cfg.use_relative_state, num_batches=1000,
+                )
+                del temp_loader, stats_dataset
+                gc.collect()
+                torch.save(rel_stats, stats_path)
+                logging.info(f"Saved relative stats to: {stats_path}")
+
+        if not is_main_process:
+            while not stats_path.exists():
+                time.sleep(5)
+
+        rel_stats = torch.load(stats_path, weights_only=True, map_location="cpu")
+
+        # Replace absolute stats with relative stats in the normalizer
+        updated_stats = dict(dataset.meta.stats)
+        updated_stats["action"] = {
+            **updated_stats["action"],
+            "mean": rel_stats["action_mean"].numpy(),
+            "std": rel_stats["action_std"].numpy(),
+        }
+        if cfg.use_relative_state and "state_mean" in rel_stats:
+            updated_stats[raw_state_key] = {
+                **updated_stats.get(raw_state_key, {}),
+                "mean": rel_stats["state_mean"].numpy(),
+                "std": rel_stats["state_std"].numpy(),
+            }
+        preprocessor = hotswap_stats(preprocessor, updated_stats)
+        logging.info("Hotswapped normalizer stats with relative stats")
+        accelerator.wait_for_everyone()
+
     step = 0  # number of policy updates (forward + backward + optim)
 
     if cfg.resume:
@@ -384,7 +444,15 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
     for _ in range(step, cfg.steps):
         start_time = time.perf_counter()
         batch = next(dl_iter)
+
+        # Convert to relative on raw data BEFORE normalization
+        if cfg.use_relative_actions:
+            batch = convert_to_relative_actions(
+                batch, state_key=raw_state_key, convert_state=cfg.use_relative_state,
+            )
+
         batch = preprocessor(batch)
+        
         train_tracker.dataloading_s = time.perf_counter() - start_time
 
         train_tracker, output_dict = update_policy(

src/lerobot/utils/relative_actions.py

@@ -0,0 +1,201 @@
+"""
+UMI-style relative actions with per-timestep normalization.
+
+Two modes supported:
+  Mode 1: Relative actions only (use_relative_state=False)
+    - Actions converted to relative, state stays absolute
+  Mode 2: Relative actions + state (use_relative_state=True, full UMI)
+    - Both actions and state converted to relative
+
+Per-timestep normalization (TRI LBM / BEHAVIOR style):
+  Training: action_norm[t] = (action_rel[t] - mean[t]) / std[t]
+  Inference: action_rel[t] = action_norm[t] * std[t] + mean[t]
+"""
+
+import torch
+from pathlib import Path
+
+
+class PerTimestepNormalizer:
+    """Per-timestep normalization using precomputed dataset statistics."""
+    
+    def __init__(self, mean: torch.Tensor, std: torch.Tensor, eps: float = 1e-8):
+        self.mean = mean
+        self.std = std
+        self.eps = eps
+        self._cache = {}  # Cache for device/dtype converted tensors
+    
+    def _get_stats(self, device, dtype):
+        """Get cached stats for device/dtype, or create and cache them."""
+        key = (device, dtype)
+        if key not in self._cache:
+            self._cache[key] = (
+                self.mean.to(device, dtype),
+                self.std.to(device, dtype),
+            )
+        return self._cache[key]
+    
+    def normalize(self, x: torch.Tensor) -> torch.Tensor:
+        mean, std = self._get_stats(x.device, x.dtype)
+        if x.dim() == 3 and mean.dim() == 2:
+            mean, std = mean.unsqueeze(0), std.unsqueeze(0)
+        return (x - mean) / (std + self.eps)
+    
+    def unnormalize(self, x: torch.Tensor) -> torch.Tensor:
+        mean, std = self._get_stats(x.device, x.dtype)
+        if x.dim() == 3 and mean.dim() == 2:
+            mean, std = mean.unsqueeze(0), std.unsqueeze(0)
+        return x * (std + self.eps) + mean
+    
+    def save(self, path: Path | str):
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        torch.save({"mean": self.mean.cpu(), "std": self.std.cpu(), "eps": self.eps}, path)
+    
+    @classmethod
+    def load(cls, path: Path | str) -> "PerTimestepNormalizer":
+        data = torch.load(path, weights_only=True, map_location="cpu")
+        return cls(data["mean"], data["std"], data.get("eps", 1e-8))
+
+
+def compute_relative_action_stats(
+    dataloader,
+    state_key: str = "observation.state",
+    num_batches: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute per-timestep mean/std from relative actions."""
+    all_rel = []
+    for i, batch in enumerate(dataloader):
+        if num_batches is not None and i >= num_batches:
+            break
+        action, state = batch["action"], batch[state_key]
+        current_pos = state[:, -1, :] if state.dim() == 3 else state
+        min_dim = min(action.shape[-1], current_pos.shape[-1])
+        rel = action.clone()
+        rel[..., :min_dim] -= current_pos[:, None, :min_dim]
+        all_rel.append(rel)
+
+    all_rel = torch.cat(all_rel, dim=0)
+    return all_rel.mean(dim=0), all_rel.std(dim=0).clamp(min=1e-6)
+
+
+def compute_global_relative_stats(
+    dataloader,
+    state_key: str = "observation.state",
+    convert_state: bool = True,
+    num_batches: int | None = None,
+) -> dict[str, torch.Tensor]:
+    """Compute global mean/std for relative actions (and state) across all timesteps.
+
+    Returns stats compatible with the standard MEAN_STD normalizer (shape = action_dim).
+    """
+    all_rel_actions = []
+    all_rel_states = []
+    for i, batch in enumerate(dataloader):
+        if num_batches is not None and i >= num_batches:
+            break
+        action, state = batch["action"], batch[state_key]
+        current_pos = state[:, -1, :] if state.dim() == 3 else state
+
+        min_dim = min(action.shape[-1], current_pos.shape[-1])
+        rel = action.clone()
+        rel[..., :min_dim] -= current_pos[:, None, :min_dim]
+        all_rel_actions.append(rel.reshape(-1, rel.shape[-1]))
+
+        if convert_state:
+            if state.dim() == 3:
+                rel_state = state - current_pos[:, None, :]
+            else:
+                rel_state = torch.zeros_like(state)
+            all_rel_states.append(rel_state.reshape(-1, rel_state.shape[-1]))
+
+    all_rel_actions = torch.cat(all_rel_actions, dim=0)
+    result = {
+        "action_mean": all_rel_actions.mean(dim=0),
+        "action_std": all_rel_actions.std(dim=0).clamp(min=1e-6),
+    }
+    if convert_state and all_rel_states:
+        all_rel_states = torch.cat(all_rel_states, dim=0)
+        result["state_mean"] = all_rel_states.mean(dim=0)
+        result["state_std"] = all_rel_states.std(dim=0).clamp(min=1e-6)
+    return result
+
+
+def convert_to_relative(
+    batch: dict,
+    state_key: str = "observation.state",
+    convert_state: bool = True,
+) -> dict:
+    """
+    Convert actions (and optionally state) to relative.
+    
+    Args:
+        batch: Training batch with "action" and state_key
+        state_key: Key for observation state
+        convert_state: If True, also convert state to relative (full UMI mode)
+    """
+    if "action" not in batch or state_key not in batch:
+        return batch
+    
+    action = batch["action"]
+    state = batch[state_key]
+    batch = batch.copy()
+    
+    # Get current position as reference
+    current_pos = state[:, -1, :] if state.dim() == 3 else state
+    
+    # Convert state if requested
+    if convert_state:
+        if state.dim() == 3:
+            batch[state_key] = state - current_pos[:, None, :]
+        else:
+            batch[state_key] = torch.zeros_like(state)
+    
+    # Convert actions to relative
+    min_dim = min(action.shape[-1], current_pos.shape[-1])
+    rel_action = action.clone()
+    rel_action[..., :min_dim] -= current_pos[:, None, :min_dim]
+    batch["action"] = rel_action
+    
+    return batch
+
+
+# Backward compatibility alias
+convert_to_relative_actions = convert_to_relative
+
+
+def convert_state_to_relative(state: torch.Tensor) -> torch.Tensor:
+    """Convert state to relative (for inference with use_relative_state=True)."""
+    if state.dim() == 1:
+        return torch.zeros_like(state)
+    current_pos = state[-1, :] if state.dim() == 2 else state[:, -1, :]
+    if state.dim() == 2:
+        return state - current_pos[None, :]
+    return state - current_pos[:, None, :]
+
+
+def convert_from_relative_actions(
+    relative_actions: torch.Tensor,
+    current_pos: torch.Tensor,
+) -> torch.Tensor:
+    """Convert relative actions back to absolute for robot execution."""
+    current_pos = current_pos.to(relative_actions.device, relative_actions.dtype)
+    min_dim = min(relative_actions.shape[-1], current_pos.shape[-1])
+    absolute = relative_actions.clone()
+    
+    if relative_actions.dim() == 2:
+        absolute[..., :min_dim] += current_pos[:min_dim]
+    elif relative_actions.dim() == 3:
+        absolute[..., :min_dim] += current_pos[None, None, :min_dim]
+    else:
+        absolute[..., :min_dim] += current_pos[:min_dim]
+    
+    return absolute
+
+
+def convert_from_relative_actions_dict(
+    relative_actions: dict[str, float],
+    current_pos: dict[str, float],
+) -> dict[str, float]:
+    """Convert relative actions back to absolute (dict version for inference)."""
+    return {k: v + current_pos.get(k, 0.0) for k, v in relative_actions.items()}