diff --git a/configs/_examples/_FULL_REFERENCE.yaml b/configs/_examples/_FULL_REFERENCE.yaml index df093d1..cccf443 100644 --- a/configs/_examples/_FULL_REFERENCE.yaml +++ b/configs/_examples/_FULL_REFERENCE.yaml @@ -109,6 +109,11 @@ data: # actions: action # prompt: prompt + # ---- LeRobotTron2DataConfig-specific ---- + # Same options and 20D layout as BiFlexiv (no optional head joints). + # Cameras: observation.images.top -> head, plus left_wrist and right_wrist. + # Example: pi05_base_tron2rt_pnp_0918.yaml + # ---- LeRobotXtacUmiDataConfig-specific (uncomment when used) ---- # Handheld bi_taccap_gripper rig. 20D state/action in the rig's own per-side- # grouped order: [left_tcp(0-8), left_gripper(9), right_tcp(10-18), right_gripper(19)] diff --git a/configs/_examples/pi05_base_tron2rt_pnp_0918.yaml b/configs/_examples/pi05_base_tron2rt_pnp_0918.yaml new file mode 100644 index 0000000..fc1811e --- /dev/null +++ b/configs/_examples/pi05_base_tron2rt_pnp_0918.yaml @@ -0,0 +1,50 @@ +# Launch template (8xH100, hb-openpi env). XLA_FLAGS must be set before Python/JAX +# starts and therefore cannot be injected by a TrainConfig YAML. The flags below measured +# -3.5% step time on 2026-09-04 (FSDP collective combining + pipelined collectives); see +# docs/training-optimization.md for the performance results and numerical validation. +# env -u LD_LIBRARY_PATH python scripts/check_cuda_stack.py +# env -u LD_LIBRARY_PATH \ +# XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true \ +# --xla_gpu_all_gather_combine_threshold_bytes=1073741824 \ +# --xla_gpu_reduce_scatter_combine_threshold_bytes=1073741824 \ +# --xla_gpu_all_reduce_combine_threshold_bytes=1073741824 \ +# --xla_gpu_enable_pipelined_all_gather=true \ +# --xla_gpu_enable_pipelined_reduce_scatter=true \ +# --xla_gpu_enable_while_loop_double_buffering=true" \ +# XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 \ +# python scripts/train.py pi05_base_tron2rt_pnp_0918 \ +# --exp-name= +model: + type: Pi0Config + # Match the H100 template; token length has not been scanned for this dataset. + max_token_len: 128 + pi05: true + discrete_state_input: true + enable_training_time_rtc: False + # Encode all three visual views in one SigLIP call. + batch_image_views: true + use_cudnn_attention: true + cudnn_attention_dtype: float16 +weight_loader: + type: CheckpointWeightLoader + params_path: gs://openpi-assets/checkpoints/pi05_base/params +ema_decay: null +data: + type: LeRobotTron2DataConfig + # Uses the local LeRobot cache under ~/.cache/huggingface/lerobot/. + repo_id: Xense/tron2rt-pnp-0918 + use_delta_cartesian_actions: true + base_config: + prompt_from_task: true + # The four tactile streams are not consumed by this policy. + tactile: false + default_prompt: Put the block into the box +batch_size: 256 +num_workers: 64 +num_train_steps: 80000 +# Whole-model grad/parameter norms are computed only on these logging steps. +log_interval: 100 +save_interval: 10000 +keep_period: 10000 +fsdp_devices: 8 +strict_batch_order: false diff --git a/examples/tron2_rt/env.py b/examples/tron2_rt/env.py new file mode 100644 index 0000000..3696673 --- /dev/null +++ b/examples/tron2_rt/env.py @@ -0,0 +1,197 @@ +"""OpenPI Environment wrapper for the LimX TRON2 dual-arm robot (Tron2RT driver). + +Matches the layout the policy was trained on (``LeRobotTron2DataConfig``, +e.g. ``Xense/tron2rt-pnp-0918``): + +State/action (20D): + [left_tcp.x/y/z/r1-r6 (0-8), right_tcp.x/y/z/r1-r6 (9-17), + left_gripper.pos (18), right_gripper.pos (19)] + +Cameras (driver name -> policy name): + top -> head + left_wrist -> left_wrist + right_wrist -> right_wrist + (tactile streams are ignored by the policy) + +Frames: ``Tron2RT.get_observation`` reports TCP poses in the tool-calibrated TCP +frame, and the recorded dataset actions are in that same frame, so the policy +outputs TCP-frame targets. ``Tron2RT.send_action`` however consumes seventh-axis +flange poses (the frame of the native IK). ``apply_action`` therefore maps each +predicted TCP pose back to the flange with ``T_flange = T_tcp @ inv(T_flange_tcp)`` +before sending. Without a calibration file the transform is identity and this +is a no-op. +""" + +import time +from typing import override + +import einops +from lerobot.robots.tron2_rt import Tron2RT +from lerobot.robots.tron2_rt import tron2_rt as _tron2_rt_driver +from lerobot.robots.tron2_rt.config_tron2_rt import Tron2RTConfig +from lerobot.robots.tron2_rt.config_tron2_rt import Tron2RTControlMode +from lerobot.utils.robot_utils import get_logger +import numpy as np +from xense_client import image_tools +from xense_client.runtime import environment as _environment + +logger = get_logger("Tron2RTEnv") + +_SIDES = ("left", "right") +_TCP_KEYS = {side: _tron2_rt_driver._tcp_keys(side) for side in _SIDES} +_GRIPPER_KEYS = ("left_gripper.pos", "right_gripper.pos") + +# Driver camera name -> policy camera name (must match BiFlexivInputs.EXPECTED_CAMERAS). +_CAMERA_MAP = { + "top": "head", + "left_wrist": "left_wrist", + "right_wrist": "right_wrist", +} + +# Action dimension labels for debug logging (20D Cartesian) +_ACTION_LABELS = [ + *(f"L.{k.split('.')[1]}" for k in _TCP_KEYS["left"]), + *(f"R.{k.split('.')[1]}" for k in _TCP_KEYS["right"]), + "L.grip", + "R.grip", +] + + +class Tron2RTEnvironment(_environment.Environment): + """OpenPI environment for the TRON2 robot in Cartesian RT mode. + + get_observation() reads state + cameras fresh; apply_action() only submits + a waypoint to the native 300 Hz publisher (which interpolates between + waypoints using their measured send times), so obs and action I/O stay + decoupled like the BiFlexiv example. + """ + + def __init__( + self, + robot_config: Tron2RTConfig, + render_height: int = 224, + render_width: int = 224, + setup_robot: bool = True, + ) -> None: + if Tron2RT is None: + raise RuntimeError( + "lerobot.robots.tron2_rt.Tron2RT is unavailable: the libtron2rt / tron2-ik runtime " + "could not be imported in this environment." + ) + if robot_config.control_mode != Tron2RTControlMode.CARTESIAN or not robot_config.observe_tcp: + raise ValueError( + "The TRON2 policy uses Cartesian TCP state/actions: need control_mode=cartesian and observe_tcp=True." + ) + # The 20D vectors end in left/right_gripper.pos, so both sides need a gripper. + missing = [ + side + for side, cfg in (("left", robot_config.left_gripper), ("right", robot_config.right_gripper)) + if cfg is None + ] + if missing: + raise ValueError( + f"No gripper configured on: {', '.join(missing)}. The policy's state and action " + "vectors are 20D ending in left/right_gripper.pos, so both sides need one." + ) + + self.config = robot_config + self.robot = Tron2RT(robot_config) + self._render_height = render_height + self._render_width = render_width + self._step_count = 0 + self._episode_count = 0 + # T_flange_tcp per side; filled in on connect. + self._inv_tool = {side: np.eye(4) for side in _SIDES} + + if setup_robot: + self.setup_robot() + + def setup_robot(self) -> None: + logger.info(f"Connecting to TRON2 RT at {self.config.robot_ip} (go_to_start={self.config.go_to_start})...") + self.robot.connect(calibrate=False, go_to_start=self.config.go_to_start) + # Use the exact transforms the driver applies to its TCP observations, so + # the TCP -> flange inverse is guaranteed to match. + self.robot._load_tool_transforms() + for side in _SIDES: + tool = np.asarray(self.robot._tool_transforms[side], dtype=np.float64) + self._inv_tool[side] = np.linalg.inv(tool) + if not np.allclose(tool, np.eye(4), atol=1e-12): + logger.info(f"{side} TCP calibration active: t_flange_tcp={tool[:3, 3].round(4).tolist()}") + logger.info("TRON2 RT connected and ready") + + @override + def reset(self) -> None: + self._episode_count += 1 + self._step_count = 0 + # connect() already moved to the start pose for the first episode. + if self._episode_count > 1 and self.config.go_to_start: + logger.info("Resetting TRON2 RT to start pose...") + self.robot.reset_to_initial_position() + logger.info("TRON2 RT reset completed") + + @override + def is_episode_complete(self) -> bool: + return False + + @override + def get_observation(self) -> dict: + raw_obs = self.robot.get_observation() + + state = np.array( + [raw_obs[k] for side in _SIDES for k in _TCP_KEYS[side]] + [raw_obs[k] for k in _GRIPPER_KEYS], + dtype=np.float32, + ) + + raw_images = {} + processed_images = {} + for src, dst in _CAMERA_MAP.items(): + img = raw_obs.get(src) + if img is None: + logger.debug(f"Camera {src} not found in observation") + continue + raw_images[dst] = img + resized = image_tools.resize_with_pad(img[None], self._render_height, self._render_width)[0] + # (H, W, C) -> (C, H, W) for OpenPI policy input + processed_images[dst] = einops.rearrange(resized, "h w c -> c h w") + + return { + "state": state, + "images": processed_images, + "images_raw": raw_images, + } + + @override + def apply_action(self, action: dict) -> None: + self._step_count += 1 + actions = action.get("actions") + if actions is None: + return + parts = " | ".join(f"{lbl}={v:+.4f}" for lbl, v in zip(_ACTION_LABELS, actions)) + logger.debug(f"Step {self._step_count}: {parts}") + self.robot.send_action(self._build_action_dict(np.asarray(actions))) + + def _build_action_dict(self, action: np.ndarray) -> dict[str, float]: + """20D TCP-frame policy action -> Tron2RT flange-frame action dict.""" + action_dict: dict[str, float] = {} + for side, offset in (("left", 0), ("right", 9)): + tcp = _tron2_rt_driver._pose9d_to_matrix(action[offset : offset + 9], side) + flange = tcp @ self._inv_tool[side] + pose9d = _tron2_rt_driver._matrix_to_pose9d(flange, side) + action_dict.update(zip(_TCP_KEYS[side], pose9d, strict=True)) + action_dict["left_gripper.pos"] = float(np.clip(action[18], 0.0, 1.0)) + action_dict["right_gripper.pos"] = float(np.clip(action[19], 0.0, 1.0)) + return action_dict + + def disconnect(self) -> None: + # Not gated on robot.is_connected: that aggregate is False as soon as a + # camera or gripper is unhealthy, which is exactly when the arm still + # needs its safe reset and RT shutdown. Tron2RT.disconnect() is a no-op + # when nothing is connected. reset_on_disconnect (default True) returns + # the robot to the start pose. + logger.info("Disconnecting TRON2 RT...") + try: + self.robot.disconnect() + time.sleep(0.5) + logger.info("TRON2 RT disconnected") + except Exception as e: + logger.warn(f"Error during disconnect: {e}") diff --git a/examples/tron2_rt/main.py b/examples/tron2_rt/main.py new file mode 100644 index 0000000..1b0312f --- /dev/null +++ b/examples/tron2_rt/main.py @@ -0,0 +1,298 @@ +#!/usr/bin/env python +"""Main script for LimX TRON2 dual-arm robot inference with OpenPI. + +Drives the lerobot ``Tron2RT`` driver (native 300 Hz joint publisher with +Cartesian waypoint streaming) against a policy served from e.g. +``configs/_examples/pi05_base_tron2rt_pnp_0918.yaml``. The bench is described by +the recipe selected by ``--args.robot-recipe`` (robot IP, Bridge camera host, +start joints, TacCap grippers); the flags below override the few knobs that +change between runs. With no recipe, ``Tron2RTConfig`` defaults are used. + +Requires the ``Tron2RT`` driver from the ``feat/taccap_web`` branch of +XenseRobotics-AI/lerobot-tron2 (the shipped recipe uses its TacCap remote +fields, which the driver's main branch does not have yet). + +--args.run picks a run YAML from runs/, which presets any of the flags below. +Flags still override the file. See examples/run_config.py. + +Example usage: + +# Serve the policy (GPU machine) +uv run scripts/serve_policy.py policy:checkpoint \\ + --policy.config=pi05_base_tron2rt_pnp_0918 --policy.dir= + +# Basic inference +python -m examples.tron2_rt.main --args.robot-recipe default \ + --args.host 192.168.2.100 --args.port 8000 + +# Dry run (robot connected but actions not sent) +python -m examples.tron2_rt.main --args.robot-recipe default \ + --args.run dry-run --args.host 192.168.2.100 + +# Override the prompt sent to the server +python -m examples.tron2_rt.main --args.host 192.168.2.100 --args.prompt "Put the block into the box" + + +python -m examples.tron2_rt.main \ + --args.robot-recipe default \ + --args.host \ + --args.port 8000 \ + --args.dry-run +""" + +from dataclasses import dataclass +import os +import pathlib +import signal +import threading +from typing import override + +from lerobot.grippers import TaccapFollowerConfig +from lerobot.robots.tron2_rt.config_tron2_rt import Tron2RTConfig +from lerobot.utils.robot_utils import get_logger +from xense_client import action_chunk_broker +from xense_client import websocket_client_policy as _websocket_client_policy +from xense_client.runtime import environment as _environment +from xense_client.runtime import runtime as _runtime +from xense_client.runtime.agents import policy_agent as _policy_agent + +import examples.run_config as _run_config +import examples.tron2_rt.env as _env +import examples.tron2_rt.recipe as _recipe + +logger = get_logger("Tron2RTMain") + +# Run YAMLs shipped with this example; --args.run resolves bare names here. +RUNS_DIR = pathlib.Path(__file__).parent / "runs" + + +class DryRunEnvironmentWrapper(_environment.Environment): + """Intercepts policy actions and prints them without executing on robot.""" + + def __init__(self, wrapped_env: _env.Tron2RTEnvironment): + self._wrapped_env = wrapped_env + self._step_count = 0 + self._episode_count = 0 + + @override + def reset(self) -> None: + self._episode_count += 1 + self._step_count = 0 + logger.info(f"\n{'=' * 80}") + logger.info(f"Episode {self._episode_count} - reset (dry run)") + logger.info(f"{'=' * 80}\n") + self._wrapped_env.reset() + + @override + def is_episode_complete(self) -> bool: + return self._wrapped_env.is_episode_complete() + + @override + def get_observation(self) -> dict: + return self._wrapped_env.get_observation() + + @override + def apply_action(self, action: dict) -> None: + self._step_count += 1 + actions = action.get("actions") + if actions is None: + return + logger.info(f"\n{'─' * 80}") + logger.info(f"Step {self._step_count} - policy action (20D Cartesian, TCP frame):") + logger.info(f"{'─' * 80}") + for i, (label, value) in enumerate(zip(_env._ACTION_LABELS, actions)): + logger.info(f" [{i:2d}] {label:8s}: {value:+.6f}") + logger.info("DRY RUN: action NOT sent to robot") + logger.info(f"{'─' * 80}\n") + + def disconnect(self) -> None: + self._wrapped_env.disconnect() + + +class PromptEnvironmentWrapper(_environment.Environment): + """Attaches a fixed language prompt to every observation.""" + + def __init__(self, wrapped_env: _environment.Environment, prompt: str): + self._wrapped_env = wrapped_env + self._prompt = prompt + + @override + def reset(self) -> None: + self._wrapped_env.reset() + + @override + def is_episode_complete(self) -> bool: + return self._wrapped_env.is_episode_complete() + + @override + def get_observation(self) -> dict: + return {**self._wrapped_env.get_observation(), "prompt": self._prompt} + + @override + def apply_action(self, action: dict) -> None: + self._wrapped_env.apply_action(action) + + +@dataclass +class Args: + """Arguments for TRON2 RT inference. + + Any of these can be preset in a run YAML under runs/ and selected with + --args.run. Flags still win over the file; see examples/run_config.py. + """ + + # Which run YAML to take the settings below from. A name resolves against + # examples/tron2_rt/runs/; a path loads any YAML. + run: str | None = None + + # Which physical bench recipe to load. A name resolves against + # examples/tron2_rt/recipes/; a path loads any recipe YAML. When omitted, + # the Tron2RTConfig defaults are used for backwards compatibility. + robot_recipe: str | None = None + + # Policy server + host: str = "localhost" + port: int = 8000 + # Language prompt sent with each observation. None = use the server's + # default_prompt (the train config sets "Put the block into the box"). + prompt: str | None = None + + # Robot (overrides on top of Tron2RTConfig defaults) + robot_ip: str | None = None + camera_host: str | None = None + go_to_start: bool = True + # Return to the start pose on exit (Ctrl+C / end of run). + reset_on_disconnect: bool = True + # The policy does not consume tactile streams; leave them off to save USB/CPU. + enable_tactile: bool = False + # Must match the TCP calibration used when recording the training data: + # observations and dataset actions are expressed in the calibrated TCP frame. + use_tool_calibration: bool = True + # Explicit tool_tcp.yaml; None = the driver's default calibration dir. + tool_calibration_path: str | None = None + # Tron2RT driver log level. + log_level: str = "INFO" + + # Image rendering + render_height: int = 224 + render_width: int = 224 + + # Runtime settings (dataset is recorded at 30 fps) + runtime_hz: float = 30.0 + num_episodes: int = 1 + max_episode_steps: int = 1000000 + + # Dry run mode + dry_run: bool = False + + # Action chunking: number of actions executed from each predicted chunk + # before re-querying the server (model action_horizon is 50). + action_horizon: int = 50 + + +def make_robot_config(args: Args) -> Tron2RTConfig: + overrides = { + "robot_ip": args.robot_ip, + "camera_host": args.camera_host, + "go_to_start": args.go_to_start, + "reset_on_disconnect": args.reset_on_disconnect, + "enable_tactile": args.enable_tactile, + "use_tool_calibration": args.use_tool_calibration, + "tool_calibration_path": pathlib.Path(args.tool_calibration_path) if args.tool_calibration_path else None, + "log_level": args.log_level, + } + if args.robot_recipe is not None: + return _recipe.load_robot_config(args.robot_recipe, **overrides) + + # Preserve the original no-recipe behavior: Tron2RTConfig supplies the + # network, startup pose, camera and TacCap defaults, while the CLI flags + # below still override the run-tunable fields. + robot_kwargs = {key: value for key, value in overrides.items() if value is not None and key != "enable_tactile"} + robot_kwargs["gripper"] = TaccapFollowerConfig(enable_tactile=args.enable_tactile) + return Tron2RTConfig(**robot_kwargs) + + +def main(args: Args) -> None: + logger.info(_run_config.describe(args, Args, RUNS_DIR)) + + if args.robot_recipe is not None: + recipe_path = _recipe.resolve_recipe_path(args.robot_recipe) + logger.info(f"Robot recipe: {recipe_path}") + + # Build (and validate) the robot config before connecting: the websocket + # client blocks until the policy server answers. + robot_config = make_robot_config(args) + logger.info( + f"TRON2 RT: robot_ip={robot_config.robot_ip}, camera_host={robot_config.camera_host}, " + f"gripper={robot_config.gripper.type if robot_config.gripper else None}, " + f"tactile={args.enable_tactile}" + ) + + ws_client_policy = _websocket_client_policy.WebsocketClientPolicy( + host=args.host, + port=args.port, + ) + logger.info(f"Server metadata: {ws_client_policy.get_server_metadata()}") + + base_environment = _env.Tron2RTEnvironment( + robot_config=robot_config, + render_height=args.render_height, + render_width=args.render_width, + setup_robot=True, + ) + + environment: _environment.Environment = base_environment + if args.dry_run: + logger.info("DRY RUN mode: actions will be printed, not executed") + environment = DryRunEnvironmentWrapper(base_environment) + if args.prompt is not None: + environment = PromptEnvironmentWrapper(environment, args.prompt) + + policy = action_chunk_broker.ActionChunkBroker( + policy=ws_client_policy, + action_horizon=args.action_horizon, + ) + + runtime = _runtime.Runtime( + environment=environment, + agent=_policy_agent.PolicyAgent(policy=policy), + subscribers=[], + max_hz=args.runtime_hz, + num_episodes=args.num_episodes, + max_episode_steps=args.max_episode_steps, + ) + + # SIGINT handling: first press asks the runtime to stop cleanly; `finally` + # then disconnects, which (reset_on_disconnect) returns the robot to the + # start pose. A second Ctrl+C escapes via os._exit. + _shutdown_in_progress = threading.Event() + + def signal_handler(sig, frame): + if _shutdown_in_progress.is_set(): + logger.warn("Second Ctrl+C — forcing exit. Robot may not return to start pose.") + os._exit(1) + _shutdown_in_progress.set() + logger.info("Ctrl+C — stopping runtime gracefully (press Ctrl+C again to force exit)") + runtime.request_stop() + + signal.signal(signal.SIGINT, signal_handler) + + try: + runtime.run() + except KeyboardInterrupt: + logger.info("Keyboard interrupt") + except Exception as e: + logger.error(f"Runtime error: {e}") + import traceback + + traceback.print_exc() + raise + finally: + try: + base_environment.disconnect() + except Exception as e: + logger.warn(f"Error disconnecting: {e}") + + +if __name__ == "__main__": + main(_run_config.cli(main, Args, RUNS_DIR)) diff --git a/examples/tron2_rt/recipe.py b/examples/tron2_rt/recipe.py new file mode 100644 index 0000000..2542635 --- /dev/null +++ b/examples/tron2_rt/recipe.py @@ -0,0 +1,35 @@ +"""Bench recipes for the TRON2 RT OpenPI example. + +The recipe contains physical bench settings (network endpoints, startup pose, +Bridge camera and gripper). Per-run policy/runtime settings stay in +``runs/`` and are applied by :mod:`examples.tron2_rt.main`. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from lerobot.robots.tron2_rt.config_tron2_rt import Tron2RTConfig + +import examples.flexiv_recipe as _shared + +RECIPES_DIR = Path(__file__).parent / "recipes" + + +def available_recipes() -> list[str]: + """Return the names of recipes shipped with this example.""" + + return _shared.available_recipes(RECIPES_DIR) + + +def resolve_recipe_path(recipe: str | Path) -> Path: + """Resolve a recipe name or an explicit YAML path.""" + + return _shared.resolve_recipe_path(RECIPES_DIR, recipe) + + +def load_robot_config(recipe: str | Path, **overrides: Any) -> Tron2RTConfig: + """Decode a recipe and apply non-``None`` run-time overrides.""" + + return _shared.load_robot_config(RECIPES_DIR, Tron2RTConfig, recipe, **overrides) diff --git a/examples/tron2_rt/recipes/default.yaml b/examples/tron2_rt/recipes/default.yaml new file mode 100644 index 0000000..32b298b --- /dev/null +++ b/examples/tron2_rt/recipes/default.yaml @@ -0,0 +1,85 @@ +# Inference recipe — LimX TRON2 dual-arm native RT controller. +# +# This file describes the physical bench. Policy-server and runtime tuning +# belongs in examples/tron2_rt/runs/ or on the command line. The values below +# mirror Tron2RTConfig's current deployment defaults; update the IPs and pose +# when this recipe is used on a different bench. +# +# python -m examples.tron2_rt.main \ +# --args.robot-recipe default --args.host +# +# Requires the Tron2RT driver from XenseRobotics-AI/lerobot-tron2 branch +# `feat/taccap_web` (not yet on main): `taccap_remote_url` and +# `taccap_remote_auto_enable` below only exist there, and draccus rejects them +# against the main-branch Tron2RTConfig. + +robot: + type: tron2_rt + + # ── LimX connection ─────────────────────────────────────────────────────── + robot_ip: 10.192.1.2 + robot_port: 5000 + camera_host: 10.192.1.4 + + # ── Control/observation schema ──────────────────────────────────────────── + # The OpenPI policy uses two Cartesian TCP poses plus two normalized grippers. + control_mode: cartesian + observe_tcp: true + observe_joints: false + use_head: false + use_grippers: true + left_use_gripper: true + right_use_gripper: true + + # ── Startup and shutdown pose ───────────────────────────────────────────── + # Arm joints are in radians, ordered left J1..J7 then right J1..J7. + init_joints: + [ + 0.026899, + 0.2612, + -0.02709991, + -1.5477003, + 0.265, + 0.0180999, + -0.0614999, + 0.008999, + -0.269, + 0.02069998, + -1.5567001, + -0.254, + -0.02309972, + 0.06469989, + ] + # Head is [pitch, yaw] in radians. It is retained for reset even though + # this policy does not expose a head action. + init_head: [1.0, 0.0] + go_to_start: true + reset_on_disconnect: true + start_move_time_s: 3.0 + + # The training data and policy actions use the calibrated TCP frame. When + # omitted, the driver looks for tool_tcp.yaml in its normal calibration dir. + use_tool_calibration: true + + # ── Bridge camera ────────────────────────────────────────────────────────── + # Tron2RTConfig supplies the fixed top/head stream: + # /camera/top/color/image_raw/compressed @ 640x480, 60 Hz + # TacCap injects left_wrist/right_wrist (and optional tactile) streams below. + # Keep `camera_host` above in sync with the Bridge machine. + + # ── TacCap gripper pair ──────────────────────────────────────────────────── + # One shared block is cloned for left/right and stamped with the physical side + # by Tron2RTConfig. The .4 service owns the USB devices and camera streams. + gripper: + type: taccap_follower + kp: 8.0 + kd: 1.0 + feedforward_torque: 0.0 + control_hz: 100 + init_open: true + require_calibrated: true + auto_discover_cameras: true + enable_tactile: false + # The top-level endpoint below is copied into both side-specific followers. + taccap_remote_url: http://10.192.1.4:8765 + taccap_remote_auto_enable: true diff --git a/examples/tron2_rt/runs/dry-run.yaml b/examples/tron2_rt/runs/dry-run.yaml new file mode 100644 index 0000000..94e7ba1 --- /dev/null +++ b/examples/tron2_rt/runs/dry-run.yaml @@ -0,0 +1,14 @@ +# Bench check: connect to TRON2, run the policy, print the actions instead of +# sending them to the robot. +# python -m examples.tron2_rt.main --args.run dry-run --args.host + +robot_recipe: default +dry_run: true + +host: localhost +port: 8000 + +runtime_hz: 30.0 +num_episodes: 1 +max_episode_steps: 200 +log_level: DEBUG diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py index 9f5b059..ca46179 100644 --- a/src/openpi/training/config.py +++ b/src/openpi/training/config.py @@ -489,6 +489,40 @@ def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig ) +@dataclasses.dataclass(frozen=True) +class LeRobotTron2DataConfig(LeRobotBiFlexivDataConfig): + """TRON2 dual-arm Cartesian datasets with grippers and no head joints. + + The 20D state/action layout matches BiFlexiv exactly: left TCP (0:9), + right TCP (9:18), then left/right grippers (18:20). Reuse its policy + transforms and delta encoding (TCP relative, grippers absolute). + + Xense/tron2rt-pnp-0918 records top, left_wrist and right_wrist cameras; + map top to the policy's head input. The plain tron2 driver's default + left/right camera names require a custom repack_transforms override. + Optional head joints and gripper-less recordings are not supported. + """ + + repack_transforms: tyro.conf.Suppress[_transforms.Group] = dataclasses.field( + default=_transforms.Group( + inputs=[ + _transforms.RepackTransform( + { + "images": { + "head": "observation.images.top", + "left_wrist": "observation.images.left_wrist", + "right_wrist": "observation.images.right_wrist", + }, + "state": "observation.state", + "actions": "action", + "prompt": "task", + } + ) + ] + ) + ) + + @dataclasses.dataclass(frozen=True) class LeRobotDobotNova5DataConfig(DataConfigFactory): """Data config for the single right-arm Dobot Nova5 DH in LeRobot format. diff --git a/src/openpi/training/registry.py b/src/openpi/training/registry.py index 169afcb..30db27d 100644 --- a/src/openpi/training/registry.py +++ b/src/openpi/training/registry.py @@ -74,6 +74,7 @@ def _populate_data_configs() -> None: "LeRobotDROIDDataConfig": _config.LeRobotDROIDDataConfig, "RLDSDroidDataConfig": _config.RLDSDroidDataConfig, "LeRobotBiFlexivDataConfig": _config.LeRobotBiFlexivDataConfig, + "LeRobotTron2DataConfig": _config.LeRobotTron2DataConfig, "LeRobotDobotNova5DataConfig": _config.LeRobotDobotNova5DataConfig, "LeRobotXtacUmiDataConfig": _config.LeRobotXtacUmiDataConfig, # SimpleDataConfig deliberately omitted: it carries lambdas (data_transforms) diff --git a/src/openpi/training/tron2_config_test.py b/src/openpi/training/tron2_config_test.py new file mode 100644 index 0000000..73aaff1 --- /dev/null +++ b/src/openpi/training/tron2_config_test.py @@ -0,0 +1,47 @@ +import dataclasses + +import numpy as np +import pytest + +from openpi.training import config as _config +import openpi.transforms as transforms + + +@pytest.mark.parametrize("use_delta", [True, False]) +def test_tron2_dataset_pipeline_roundtrip(monkeypatch, tmp_path, use_delta): + # Keep this test offline; tokenizer/model transforms are shared and unchanged. + monkeypatch.setattr(_config.ModelTransformFactory, "__call__", lambda self, model: transforms.Group()) + config = _config.get_config("pi05_base_tron2rt_pnp_0918") + assert isinstance(config.data, _config.LeRobotTron2DataConfig) + assert config.data.repo_id == "Xense/tron2rt-pnp-0918" + factory = dataclasses.replace(config.data, use_delta_cartesian_actions=use_delta) + data_config = factory.create(tmp_path, config.model) + assert data_config.prompt_from_task + assert not data_config.tactile + assert data_config.action_sequence_keys == ("action",) + + state = np.arange(20, dtype=np.float32) + actions = np.stack([state + 1, state + 2]) + sample = { + "observation.state": state, + "action": actions.copy(), + "task": "Pick and place.", + **{ + f"observation.images.{camera}": np.full((3, 8, 8), value, dtype=np.uint8) + for value, camera in enumerate(("top", "left_wrist", "right_wrist"), start=1) + }, + } + result = transforms.compose([*data_config.repack_transforms.inputs, *data_config.data_transforms.inputs])(sample) + for value, slot in enumerate(("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb"), start=1): + np.testing.assert_array_equal(result["image"][slot], np.full((8, 8, 3), value, dtype=np.uint8)) + assert result["image_mask"][slot] + assert result["prompt"] == sample["task"] + np.testing.assert_array_equal(result["state"], state) + expected = actions.copy() + if use_delta: + expected[:, :18] -= state[:18] + np.testing.assert_array_equal(result["actions"], expected) + + result["actions"] = np.pad(result["actions"], ((0, 0), (0, 12))) + restored = transforms.compose(data_config.data_transforms.outputs)(result) + np.testing.assert_allclose(restored["actions"], actions)