Compare commits

...

11 Commits

Author SHA1 Message Date
Pepijn
1dbb8f0fc3 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
2026-04-02 13:16:06 +02:00
pepijn
936187cd76 fix relative actions: convert before normalization, use global stats
The previous implementation had a double-normalization bug: the
preprocessor normalized actions with absolute stats, then
convert_to_relative subtracted normalized state (wrong), then the
per-timestep normalizer re-normalized.

Now the correct flow is:
1. Convert batch to relative on raw data (before preprocessing)
2. Compute global relative stats (mean/std across all timesteps)
3. Hotswap the preprocessor normalizer to use relative stats
4. Preprocessor normalizes relative values correctly

This brings loss from ~3000+ down to ~0.5, matching the main branch.

Made-with: Cursor
2026-04-01 20:45:35 +00:00
pepijn
900f1a42e9 fix relative stats computation with rename_map
When rename_map maps a dataset key to observation.state, the raw
dataset used for stats computation still has the original key.
Reverse the rename_map to find the correct key.

Made-with: Cursor
2026-04-01 19:12:14 +00:00
Pepijn
5a15a6a911 use seperate process for stats computation 2026-01-07 16:52:15 +01:00
Pepijn
c23472e376 only main saves stat file 2026-01-07 15:30:25 +01:00
Pepijn
63619619bf fix data loader issue 2026-01-07 10:03:56 +01:00
Pepijn
574081ac02 fix mem bug 2026-01-03 11:34:31 +01:00
Pepijn
c5f66edff9 shuffle false 2026-01-02 22:34:57 +01:00
Pepijn
7f16e8cb09 fix 2026-01-02 19:56:42 +01:00
Pepijn
0367955590 add code for relative actions and state and unifing tasks 2026-01-02 18:58:47 +01:00
Pepijn
01c7c74070 Add relative position UMI style 2026-01-02 15:57:39 +01:00
29 changed files with 1525 additions and 7 deletions

View File

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

59
scripts/unify_tasks.py Normal file
View File

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

View File

@@ -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

View File

@@ -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}")

File diff suppressed because one or more lines are too long

View File

@@ -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>

View File

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

View File

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