ydy0615/lerobot-clone
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-06-04 21:01:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(umi): add EE replay viewer, URDF meshes, and evaluate script updates

Browse Source 
- Add replay.py script and replay_viewer.html for browser-based EE
  trajectory visualization from glannuzel/grabette-dataset
- Add viewer.html for interactive URDF inspection
- Move OpenArm URDF and meshes into openarm_follower/urdf/
- Add virtual EE target frame (openarm_right_ee_target) at 7cm from link7
- Adapt evaluate.py for single right-arm OpenArm with wrist camera
- Update docs with replay viewer usage
- Update openarm_follower config, driver, and kinematic processor

Made-with: Cursor
This commit is contained in:
Pepijn
2026-04-02 14:25:24 +02:00
parent b08a62af89
commit e627d6442e
33 changed files with 1867 additions and 203 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.gitignore vendored
View File

@@ -173,7 +173,5 @@ outputs/
# Dev folders
.cache/*
*.stl
*.urdf
*.xml
*.part

53
docs/source/umi_pi0_relative_ee.mdx
View File

@@ -110,6 +110,59 @@ The inference flow uses pi0's built-in processor pipeline — no custom wrappers
4. **Postprocessor** — `UnnormalizerProcessorStep` unnormalizes, then `AbsoluteActionsProcessorStep` adds the cached state back to get absolute EE targets.
5. **IK → Robot** — `InverseKinematicsEEToJoints` converts absolute EE targets to joint commands.
## Replay Viewer
Before running on hardware, you can visualize any dataset episode in a browser-based 3D viewer. The viewer shows the EE trajectory overlaid on the OpenArm URDF model, making it easy to sanity-check recorded data or debug unexpected behavior.
### Quick start
```bash
python examples/umi_pi0_relative_ee/replay.py
```
This extracts the trajectory from episode 0 of the default dataset, starts a local HTTP server, and opens the viewer at [http://localhost:8765/replay_viewer.html](http://localhost:8765/replay_viewer.html).
### Options
| Flag | Default | Description |
| ----------- | ---------------------------- | ------------------------------------ |
| `--repo-id` | `glannuzel/grabette-dataset` | HuggingFace dataset repo to load |
| `--episode` | `0` | Episode index to replay |
| `--port` | `8765` | HTTP server port |
| `--force` | off | Re-extract trajectory even if cached |
Example with a different dataset and episode:
```bash
python examples/umi_pi0_relative_ee/replay.py \
--repo-id myuser/my-dataset \
--episode 3 \
--port 8766
```
### Viewer controls
The panel in the top-left corner shows live EE coordinates (x, y, z, ax, ay, az) and gripper state for the current frame. Below that are transport controls:
- **Play / Pause** — toggle automatic playback.
- **Step buttons** (◀ ▶) — advance or rewind one frame at a time.
- **Reset** (⟳) — jump back to frame 0.
- **Scrubber** — drag to seek to any frame.
- **Speed selector** — 0.25×, 0.5×, 1×, 2×, or 4× playback speed.
The 3D scene uses orbit controls — click and drag to rotate, scroll to zoom, right-click drag to pan.
### Color legend
| Color | Meaning |
| ------------------ | --------------------------------------------- |
| Red sphere | Current EE position |
| Yellow trail | Past trajectory (frames already visited) |
| Dark trail | Future trajectory (frames ahead) |
| Orange ring + axes | URDF `ee_target` frame (zero-joint reference) |
The trajectory is automatically re-centered so that frame 0 aligns with the robot's `openarm_right_ee_target` link in the zero-joint pose.
## How the Pieces Fit Together
```

38
examples/umi_pi0_relative_ee/evaluate.py
View File

@@ -86,11 +86,12 @@ HF_MODEL_ID = "pepijn223/grabette-umi-pi0"
# E.g. at 10Hz with ~200ms total system latency: ceil(200 / 100) = 2.
LATENCY_SKIP_STEPS = 0
# EE feature keys produced by ForwardKinematicsJointsToEE
EE_KEYS = ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]
# EE feature keys produced by ForwardKinematicsJointsToEE (arm pose only).
# Gripper joints use absolute position control, not EE-relative.
EE_KEYS = ["x", "y", "z", "wx", "wy", "wz"]
URDF_PATH = "./SO101/so101_new_calib.urdf"
URDF_EE_FRAME = "gripper_frame_link"
URDF_PATH = "src/lerobot/robots/openarm_follower/urdf/openarm_bimanual_pybullet.urdf"
URDF_EE_FRAME = "openarm_right_link7"
def main():
@@ -101,38 +102,47 @@ def main():
side="right",
cameras=camera_config,
max_relative_target=8.0,
gripper_port="/dev/ttyUSB0",
)
robot = OpenArmFollower(robot_config)
policy = PI0Policy.from_pretrained(HF_MODEL_ID)
policy.config.latency_skip_steps = LATENCY_SKIP_STEPS
motor_names = list(robot.bus.motors.keys())
arm_motor_names = list(robot.bus.motors.keys())
gripper_motor_names = list(robot.gripper_bus.motors.keys())
kinematics_solver = RobotKinematics(
urdf_path=URDF_PATH,
target_frame_name=URDF_EE_FRAME,
joint_names=motor_names,
joint_names=arm_motor_names,
)
# FK: joint observation → EE observation (produces observation.state)
# The policy starts from the robot's current EE pose (via FK below).
# Relative actions are predicted as deltas from that pose, so no manual
# re-centering is needed — the starting point is always the live EE tip.
# FK: joint observation → EE observation (produces observation.state).
# gripper_names=[] means proximal/distal pass through as absolute positions.
robot_joints_to_ee_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
steps=[
ForwardKinematicsJointsToEE(
kinematics=kinematics_solver,
motor_names=motor_names,
motor_names=arm_motor_names,
gripper_names=[],
)
],
to_transition=observation_to_transition,
to_output=transition_to_observation,
)
# IK: EE action → joint targets
# IK: EE action → joint targets. Gripper actions are absolute and pass through.
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
steps=[
InverseKinematicsEEToJoints(
kinematics=kinematics_solver,
motor_names=motor_names,
motor_names=arm_motor_names,
gripper_names=[],
initial_guess_current_joints=True,
),
],
@@ -162,7 +172,13 @@ def main():
aggregate_pipeline_dataset_features(
pipeline=make_default_teleop_action_processor(),
initial_features=create_initial_features(
action={f"ee.{k}": PolicyFeature(type=FeatureType.ACTION, shape=(1,)) for k in EE_KEYS}
action={
**{f"ee.{k}": PolicyFeature(type=FeatureType.ACTION, shape=(1,)) for k in EE_KEYS},
**{
f"{g}.pos": PolicyFeature(type=FeatureType.ACTION, shape=(1,))
for g in gripper_motor_names
},
}
),
use_videos=True,
),

113
examples/umi_pi0_relative_ee/replay.py Normal file
View File

@@ -0,0 +1,113 @@
#!/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.
"""
Replay a dataset episode in EE frame using a browser-based URDF viewer.
Extracts ``observation.pose`` from the dataset, saves a trajectory JSON file,
then launches a local HTTP server and opens the replay viewer. The trajectory
is re-centered so frame 0 starts at the OpenArm ``openarm_right_ee_target``
EE tip (zero-joint pose).
Usage:
python replay.py
python replay.py --episode 3 --repo-id myuser/mydata
"""
from __future__ import annotations
import argparse
import http.server
import json
import os
import threading
import webbrowser
from pathlib import Path
VIEWER_DIR = Path(__file__).resolve().parents[2] / "src/lerobot/robots/openarm_follower/urdf"
TRAJECTORY_FILENAME = "trajectory_ep0.json"
def extract_trajectory(repo_id: str, episode: int, output_path: Path) -> dict:
from lerobot.datasets.lerobot_dataset import LeRobotDataset
dataset = LeRobotDataset(repo_id, episodes=[episode])
poses = dataset.select_columns("observation.pose")
actions = dataset.select_columns("action")
frames = []
for i in range(dataset.num_frames):
p = poses[i]["observation.pose"]
a = actions[i]["action"]
frames.append(
{
"x": float(p[0]),
"y": float(p[1]),
"z": float(p[2]),
"ax": float(p[3]),
"ay": float(p[4]),
"az": float(p[5]),
"proximal": float(a[0]),
"distal": float(a[1]),
}
)
payload = {"fps": dataset.fps, "num_frames": dataset.num_frames, "frames": frames}
with open(output_path, "w") as f:
json.dump(payload, f)
print(f"Extracted {dataset.num_frames} frames at {dataset.fps} FPS → {output_path}")
return payload
# ---------------------------------------------------------------------------
# Viewer mode
# ---------------------------------------------------------------------------
def serve_and_open(directory: Path, port: int = 8765):
os.chdir(directory)
handler = http.server.SimpleHTTPRequestHandler
httpd = http.server.HTTPServer(("", port), handler)
url = f"http://localhost:{port}/replay_viewer.html"
print(f"Serving at {url}")
threading.Thread(target=lambda: webbrowser.open(url), daemon=True).start()
try:
httpd.serve_forever()
except KeyboardInterrupt:
print("\nServer stopped.")
httpd.server_close()
def run_viewer(args):
trajectory_path = VIEWER_DIR / TRAJECTORY_FILENAME
if not trajectory_path.exists() or args.force:
extract_trajectory(args.repo_id, args.episode, trajectory_path)
else:
print(f"Using cached trajectory at {trajectory_path} (pass --force to re-extract)")
serve_and_open(VIEWER_DIR, args.port)
def main():
parser = argparse.ArgumentParser(description="Replay a dataset episode in EE frame (URDF viewer)")
parser.add_argument("--repo-id", default="glannuzel/grabette-dataset")
parser.add_argument("--episode", type=int, default=0)
parser.add_argument("--port", type=int, default=8765)
parser.add_argument("--force", action="store_true", help="Re-extract trajectory even if cached")
args = parser.parse_args()
run_viewer(args)
if __name__ == "__main__":
main()

3
pyproject.toml
View File

@@ -306,7 +306,8 @@ default.extend-ignore-identifiers-re = [
"thw",
"inpt",
"ROBOTIS",
"OT_VALUE"
"OT_VALUE",
"metalness",
]
# TODO: Uncomment when ready to use

4
src/lerobot/robots/bi_openarm_follower/bi_openarm_follower.py
View File

@@ -62,6 +62,8 @@ class BiOpenArmFollower(Robot):
can_bitrate=config.left_arm_config.can_bitrate,
can_data_bitrate=config.left_arm_config.can_data_bitrate,
motor_config=config.left_arm_config.motor_config,
gripper_port=config.left_arm_config.gripper_port,
gripper_motor_ids=config.left_arm_config.gripper_motor_ids,
position_kd=config.left_arm_config.position_kd,
position_kp=config.left_arm_config.position_kp,
joint_limits=config.left_arm_config.joint_limits,
@@ -80,6 +82,8 @@ class BiOpenArmFollower(Robot):
can_bitrate=config.right_arm_config.can_bitrate,
can_data_bitrate=config.right_arm_config.can_data_bitrate,
motor_config=config.right_arm_config.motor_config,
gripper_port=config.right_arm_config.gripper_port,
gripper_motor_ids=config.right_arm_config.gripper_motor_ids,
position_kd=config.right_arm_config.position_kd,
position_kp=config.right_arm_config.position_kp,
joint_limits=config.right_arm_config.joint_limits,

33
src/lerobot/robots/openarm_follower/config_openarm_follower.py
View File

@@ -28,7 +28,8 @@ LEFT_DEFAULT_JOINTS_LIMITS: dict[str, tuple[float, float]] = {
"joint_5": (-85.0, 85.0),
"joint_6": (-40.0, 40.0),
"joint_7": (-80.0, 80.0),
"gripper": (-65.0, 0.0),
"proximal": (0.0, 100.0),
"distal": (0.0, 100.0),
}
RIGHT_DEFAULT_JOINTS_LIMITS: dict[str, tuple[float, float]] = {
@@ -39,7 +40,8 @@ RIGHT_DEFAULT_JOINTS_LIMITS: dict[str, tuple[float, float]] = {
"joint_5": (-85.0, 85.0),
"joint_6": (-40.0, 40.0),
"joint_7": (-80.0, 80.0),
"gripper": (-65.0, 0.0),
"proximal": (0.0, 100.0),
"distal": (0.0, 100.0),
}
@@ -73,13 +75,8 @@ class OpenArmFollowerConfigBase:
# Camera configurations
cameras: dict[str, CameraConfig] = field(default_factory=dict)
# Motor configuration for OpenArms (7 DOF per arm)
# Arm motor configuration (7 DOF, Damiao on CAN bus)
# Maps motor names to (send_can_id, recv_can_id, motor_type)
# Based on: https://docs.openarm.dev/software/setup/configure-test
# OpenArms uses 4 types of motors:
# - DM8009 (DM-J8009P-2EC) for shoulders (high torque)
# - DM4340P and DM4340 for shoulder rotation and elbow
# - DM4310 (DM-J4310-2EC V1.1) for wrist and gripper
motor_config: dict[str, tuple[int, int, str]] = field(
default_factory=lambda: {
"joint_1": (0x01, 0x11, "dm8009"), # J1 - Shoulder pan (DM8009)
@@ -89,19 +86,18 @@ class OpenArmFollowerConfigBase:
"joint_5": (0x05, 0x15, "dm4310"), # J5 - Wrist roll (DM4310)
"joint_6": (0x06, 0x16, "dm4310"), # J6 - Wrist pitch (DM4310)
"joint_7": (0x07, 0x17, "dm4310"), # J7 - Wrist rotation (DM4310)
"gripper": (0x08, 0x18, "dm4310"), # J8 - Gripper (DM4310)
}
)
# MIT control parameters for position control (used in send_action)
# List of 8 values: [joint_1, joint_2, joint_3, joint_4, joint_5, joint_6, joint_7, gripper]
position_kp: list[float] = field(
default_factory=lambda: [240.0, 240.0, 240.0, 240.0, 24.0, 31.0, 25.0, 25.0]
)
position_kd: list[float] = field(default_factory=lambda: [5.0, 5.0, 3.0, 5.0, 0.3, 0.3, 0.3, 0.3])
# UMI-style gripper (Feetech STS3215 on serial bus)
gripper_port: str = "/dev/ttyUSB0"
gripper_motor_ids: dict[str, int] = field(default_factory=lambda: {"proximal": 1, "distal": 2})
# Values for joint limits. Can be overridden via CLI (for custom values) or by setting config.side to either 'left' or 'right'.
# If config.side is left set to None and no CLI values are passed, the default joint limit values are small for safety.
# MIT control parameters for the 7 arm joints
position_kp: list[float] = field(default_factory=lambda: [240.0, 240.0, 240.0, 240.0, 24.0, 31.0, 25.0])
position_kd: list[float] = field(default_factory=lambda: [5.0, 5.0, 3.0, 5.0, 0.3, 0.3, 0.3])
# Joint limits. Can be overridden via CLI or by setting config.side to 'left' or 'right'.
joint_limits: dict[str, tuple[float, float]] = field(
default_factory=lambda: {
"joint_1": (-5.0, 5.0),
@@ -111,7 +107,8 @@ class OpenArmFollowerConfigBase:
"joint_5": (-5.0, 5.0),
"joint_6": (-5.0, 5.0),
"joint_7": (-5.0, 5.0),
"gripper": (-5.0, 0.0),
"proximal": (0.0, 100.0),
"distal": (0.0, 100.0),
}
)

232
src/lerobot/robots/openarm_follower/openarm_follower.py
View File

@@ -22,6 +22,7 @@ from typing import Any
from lerobot.cameras.utils import make_cameras_from_configs
from lerobot.motors import Motor, MotorCalibration, MotorNormMode
from lerobot.motors.damiao import DamiaoMotorsBus
from lerobot.motors.feetech import FeetechMotorsBus, OperatingMode
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
@@ -38,8 +39,7 @@ logger = logging.getLogger(__name__)
class OpenArmFollower(Robot):
"""
OpenArms Follower Robot which uses CAN bus communication to control 7 DOF arm with a gripper.
The arm uses Damiao motors in MIT control mode.
OpenArms Follower Robot: 7 DOF Damiao arm (CAN) + UMI-style Feetech gripper (serial).
"""
config_class = OpenArmFollowerConfig
@@ -49,19 +49,17 @@ class OpenArmFollower(Robot):
super().__init__(config)
self.config = config
# Arm motors
motors: dict[str, Motor] = {}
# Arm motors (Damiao on CAN bus)
arm_motors: dict[str, Motor] = {}
for motor_name, (send_id, recv_id, motor_type_str) in config.motor_config.items():
motor = Motor(
send_id, motor_type_str, MotorNormMode.DEGREES
) # Always use degrees for Damiao motors
motor = Motor(send_id, motor_type_str, MotorNormMode.DEGREES)
motor.recv_id = recv_id
motor.motor_type_str = motor_type_str
motors[motor_name] = motor
arm_motors[motor_name] = motor
self.bus = DamiaoMotorsBus(
port=self.config.port,
motors=motors,
motors=arm_motors,
calibration=self.calibration,
can_interface=self.config.can_interface,
use_can_fd=self.config.use_can_fd,
@@ -69,6 +67,17 @@ class OpenArmFollower(Robot):
data_bitrate=self.config.can_data_bitrate if self.config.use_can_fd else None,
)
# Gripper motors (Feetech STS3215 on serial bus)
gripper_motors: dict[str, Motor] = {
name: Motor(motor_id, "sts3215", MotorNormMode.RANGE_0_100)
for name, motor_id in config.gripper_motor_ids.items()
}
self.gripper_bus = FeetechMotorsBus(
port=config.gripper_port,
motors=gripper_motors,
calibration=self.calibration,
)
if config.side is not None:
if config.side == "left":
config.joint_limits = LEFT_DEFAULT_JOINTS_LIMITS
@@ -84,7 +93,6 @@ class OpenArmFollower(Robot):
)
logger.info(f"Values used for joint limits: {config.joint_limits}.")
# Initialize cameras
self.cameras = make_cameras_from_configs(config.cameras)
@property
@@ -93,8 +101,10 @@ class OpenArmFollower(Robot):
features: dict[str, type] = {}
for motor in self.bus.motors:
features[f"{motor}.pos"] = float
features[f"{motor}.vel"] = float # Add this
features[f"{motor}.torque"] = float # Add this
features[f"{motor}.vel"] = float
features[f"{motor}.torque"] = float
for motor in self.gripper_bus.motors:
features[f"{motor}.pos"] = float
return features
@property
@@ -116,8 +126,11 @@ class OpenArmFollower(Robot):
@property
def is_connected(self) -> bool:
"""Check if robot is connected."""
return self.bus.is_connected and all(cam.is_connected for cam in self.cameras.values())
return (
self.bus.is_connected
and self.gripper_bus.is_connected
and all(cam.is_connected for cam in self.cameras.values())
)
@check_if_already_connected
def connect(self, calibrate: bool = True) -> None:
@@ -127,12 +140,12 @@ class OpenArmFollower(Robot):
We assume that at connection time, the arms are in a safe rest position,
and torque can be safely disabled to run calibration if needed.
"""
# Connect to CAN bus
logger.info(f"Connecting arm on {self.config.port}...")
self.bus.connect()
# Run calibration if needed
logger.info(f"Connecting gripper on {self.config.gripper_port}...")
self.gripper_bus.connect()
if not self.is_calibrated and calibrate:
logger.info(
"Mismatch between calibration values in the motor and the calibration file or no calibration file found"
@@ -144,7 +157,7 @@ class OpenArmFollower(Robot):
self.configure()
if self.is_calibrated:
if self.bus.is_calibrated:
self.bus.set_zero_position()
self.bus.enable_torque()
@@ -153,47 +166,39 @@ class OpenArmFollower(Robot):
@property
def is_calibrated(self) -> bool:
"""Check if robot is calibrated."""
return self.bus.is_calibrated
return self.bus.is_calibrated and self.gripper_bus.is_calibrated
def calibrate(self) -> None:
"""
Run calibration procedure for OpenArms robot.
Run calibration for both the Damiao arm and Feetech gripper.
The calibration procedure:
1. Disable torque
2. Ask user to position arms in hanging position with grippers closed
3. Set this as zero position
4. Record range of motion for each joint
5. Save calibration
Arm calibration: set zero position with arm hanging, ±90° default range.
Gripper calibration: SO100-style half-turn homing + range recording.
"""
if self.calibration:
# Calibration file exists, ask user whether to use it or run new calibration
user_input = input(
f"Press ENTER to use provided calibration file associated with the id {self.id}, or type 'c' and press ENTER to run calibration: "
)
if user_input.strip().lower() != "c":
logger.info(f"Writing calibration file associated with the id {self.id} to the motors")
self.bus.write_calibration(self.calibration)
self.gripper_bus.write_calibration(self.calibration)
return
logger.info(f"\nRunning calibration for {self}")
self.bus.disable_torque()
# Step 1: Set zero position
# --- Arm calibration (Damiao) ---
self.bus.disable_torque()
input(
"\nCalibration: Set Zero Position)\n"
"\nCalibration: Set Zero Position\n"
"Position the arm in the following configuration:\n"
" - Arm hanging straight down\n"
" - Gripper closed\n"
"Press ENTER when ready..."
)
# Set current position as zero for all motors
self.bus.set_zero_position()
logger.info("Arm zero position set.")
logger.info("Setting range: -90° to +90° for safety by default for all joints")
for motor_name, motor in self.bus.motors.items():
self.calibration[motor_name] = MotorCalibration(
id=motor.id,
@@ -202,17 +207,52 @@ class OpenArmFollower(Robot):
range_min=-90,
range_max=90,
)
self.bus.write_calibration(self.calibration)
# --- Gripper calibration (Feetech) ---
self.gripper_bus.disable_torque()
for motor in self.gripper_bus.motors:
self.gripper_bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
input("Move gripper to the middle of its range of motion and press ENTER....")
homing_offsets = self.gripper_bus.set_half_turn_homings()
gripper_motor_names = list(self.gripper_bus.motors.keys())
print(
f"Move gripper joints ({', '.join(gripper_motor_names)}) through their "
"entire ranges of motion.\nRecording positions. Press ENTER to stop..."
)
range_mins, range_maxes = self.gripper_bus.record_ranges_of_motion(gripper_motor_names)
for motor_name, m in self.gripper_bus.motors.items():
self.calibration[motor_name] = MotorCalibration(
id=m.id,
drive_mode=0,
homing_offset=homing_offsets[motor_name],
range_min=range_mins[motor_name],
range_max=range_maxes[motor_name],
)
self.gripper_bus.write_calibration(self.calibration)
self._save_calibration()
print(f"Calibration saved to {self.calibration_fpath}")
def configure(self) -> None:
"""Configure motors with appropriate settings."""
# TODO(Steven, Pepijn): Slightly different from what it is happening in the leader
"""Configure both arm (Damiao) and gripper (Feetech) motors."""
with self.bus.torque_disabled():
self.bus.configure_motors()
with self.gripper_bus.torque_disabled():
self.gripper_bus.configure_motors()
for motor in self.gripper_bus.motors:
self.gripper_bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
self.gripper_bus.write("P_Coefficient", motor, 16)
self.gripper_bus.write("I_Coefficient", motor, 0)
self.gripper_bus.write("D_Coefficient", motor, 32)
self.gripper_bus.write("Max_Torque_Limit", motor, 500)
self.gripper_bus.write("Protection_Current", motor, 250)
self.gripper_bus.write("Overload_Torque", motor, 25)
def setup_motors(self) -> None:
raise NotImplementedError(
"Motor ID configuration is typically done via manufacturer tools for CAN motors."
@@ -220,25 +260,23 @@ class OpenArmFollower(Robot):
@check_if_not_connected
def get_observation(self) -> RobotObservation:
"""
Get current observation from robot including position, velocity, and torque.
Reads all motor states (pos/vel/torque) in one CAN refresh cycle
instead of 3 separate reads.
"""
"""Read all motor states from arm (CAN) and gripper (serial), plus cameras."""
start = time.perf_counter()
obs_dict: dict[str, Any] = {}
# Arm motors (Damiao) — pos/vel/torque in one CAN refresh cycle
states = self.bus.sync_read_all_states()
for motor in self.bus.motors:
state = states.get(motor, {})
obs_dict[f"{motor}.pos"] = state.get("position", 0.0)
obs_dict[f"{motor}.vel"] = state.get("velocity", 0.0)
obs_dict[f"{motor}.torque"] = state.get("torque", 0.0)
# Capture images from cameras
# Gripper motors (Feetech) — position only
gripper_positions = self.gripper_bus.sync_read("Present_Position")
for motor, val in gripper_positions.items():
obs_dict[f"{motor}.pos"] = val
for cam_key, cam in self.cameras.items():
start = time.perf_counter()
obs_dict[cam_key] = cam.read_latest()
@@ -258,86 +296,76 @@ class OpenArmFollower(Robot):
custom_kd: dict[str, float] | None = None,
) -> RobotAction:
"""
Send action command to robot.
The action magnitude may be clipped based on safety limits.
Send action command to robot. Arm joints go to Damiao CAN bus,
gripper joints go to Feetech serial bus.
Args:
action: Dictionary with motor positions (e.g., "joint_1.pos", "joint_2.pos")
custom_kp: Optional custom kp gains per motor (e.g., {"joint_1": 120.0, "joint_2": 150.0})
custom_kd: Optional custom kd gains per motor (e.g., {"joint_1": 1.5, "joint_2": 2.0})
action: Dictionary with motor positions (e.g., "joint_1.pos", "proximal.pos")
custom_kp: Optional custom kp gains per arm motor
custom_kd: Optional custom kd gains per arm motor
Returns:
The action actually sent (potentially clipped)
"""
goal_pos = {key.removesuffix(".pos"): val for key, val in action.items() if key.endswith(".pos")}
# Apply joint limit clipping to arm
# Apply joint limit clipping
for motor_name, position in goal_pos.items():
if motor_name in self.config.joint_limits:
min_limit, max_limit = self.config.joint_limits[motor_name]
clipped_position = max(min_limit, min(max_limit, position))
if clipped_position != position:
logger.debug(f"Clipped {motor_name} from {position:.2f}° to {clipped_position:.2f}°")
logger.debug(f"Clipped {motor_name} from {position:.2f} to {clipped_position:.2f}")
goal_pos[motor_name] = clipped_position
# Cap goal position when too far away from present position.
# /!\ Slower fps expected due to reading from the follower.
if self.config.max_relative_target is not None:
# Split into arm and gripper actions
arm_motors = set(self.bus.motors.keys())
gripper_motors = set(self.gripper_bus.motors.keys())
arm_goal = {k: v for k, v in goal_pos.items() if k in arm_motors}
gripper_goal = {k: v for k, v in goal_pos.items() if k in gripper_motors}
# Cap arm goal position when too far away from present position
if self.config.max_relative_target is not None and arm_goal:
present_pos = self.bus.sync_read("Present_Position")
goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()}
goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)
goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in arm_goal.items()}
arm_goal = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)
# TODO(Steven, Pepijn): Refactor writing
# Motor name to index mapping for gains
motor_index = {
"joint_1": 0,
"joint_2": 1,
"joint_3": 2,
"joint_4": 3,
"joint_5": 4,
"joint_6": 5,
"joint_7": 6,
"gripper": 7,
}
# Arm: batch MIT control (Damiao)
if arm_goal:
arm_motor_names = list(self.bus.motors.keys())
commands = {}
for motor_name, position_degrees in arm_goal.items():
idx = arm_motor_names.index(motor_name) if motor_name in arm_motor_names else 0
if custom_kp is not None and motor_name in custom_kp:
kp = custom_kp[motor_name]
else:
kp = (
self.config.position_kp[idx]
if isinstance(self.config.position_kp, list)
else self.config.position_kp
)
if custom_kd is not None and motor_name in custom_kd:
kd = custom_kd[motor_name]
else:
kd = (
self.config.position_kd[idx]
if isinstance(self.config.position_kd, list)
else self.config.position_kd
)
commands[motor_name] = (kp, kd, position_degrees, 0.0, 0.0)
self.bus._mit_control_batch(commands)
# Use batch MIT control for arm (sends all commands, then collects responses)
commands = {}
for motor_name, position_degrees in goal_pos.items():
idx = motor_index.get(motor_name, 0)
# Use custom gains if provided, otherwise use config defaults
if custom_kp is not None and motor_name in custom_kp:
kp = custom_kp[motor_name]
else:
kp = (
self.config.position_kp[idx]
if isinstance(self.config.position_kp, list)
else self.config.position_kp
)
if custom_kd is not None and motor_name in custom_kd:
kd = custom_kd[motor_name]
else:
kd = (
self.config.position_kd[idx]
if isinstance(self.config.position_kd, list)
else self.config.position_kd
)
commands[motor_name] = (kp, kd, position_degrees, 0.0, 0.0)
self.bus._mit_control_batch(commands)
# Gripper: position control (Feetech)
if gripper_goal:
self.gripper_bus.sync_write("Goal_Position", gripper_goal)
goal_pos.update(arm_goal)
return {f"{motor}.pos": val for motor, val in goal_pos.items()}
@check_if_not_connected
def disconnect(self):
"""Disconnect from robot."""
# Disconnect CAN bus
self.bus.disconnect(self.config.disable_torque_on_disconnect)
# Disconnect cameras
self.gripper_bus.disconnect(self.config.disable_torque_on_disconnect)
for cam in self.cameras.values():
cam.disconnect()
logger.info(f"{self} disconnected.")

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/collision/link0_symp.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/collision/link1_symp.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/collision/link2_symp.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/collision/link3_symp.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/collision/link4_symp.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/collision/link5_symp.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/collision/link6_symp.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/collision/link7_symp.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/visual/link0.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/visual/link1.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/visual/link2.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/visual/link3.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/visual/link4.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/visual/link5.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/visual/link6.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/arm/v10/visual/link7.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/body/v10/collision/body_link0_symp.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/body/v10/visual/body_link0.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/ee/openarm_hand/collision/finger.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/ee/openarm_hand/collision/hand.stl LFS Normal file
View File

Binary file not shown.

BIN
src/lerobot/robots/openarm_follower/urdf/meshes/ee/openarm_hand/visual/finger.stl LFS Normal file
View File

Binary file not shown.

630
src/lerobot/robots/openarm_follower/urdf/openarm_bimanual_pybullet.urdf Normal file
View File

@@ -0,0 +1,630 @@
<?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>
<link name="openarm_right_ee_target">
<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_ee_target_joint" type="fixed">
<parent link="openarm_right_link7" />
<child link="openarm_right_ee_target" />
<origin rpy="0 0 0" xyz="0 0.0 0.07" />
</joint>
<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>

408
src/lerobot/robots/openarm_follower/urdf/replay_viewer.html Normal file
View File

@@ -0,0 +1,408 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8" />
<title>Dataset Replay — EE Frame Viewer</title>
<style>
* { margin: 0; padding: 0; box-sizing: border-box; }
body { background: #0d1117; overflow: hidden; font-family: 'JetBrains Mono', monospace; color: #c9d1d9; }
canvas { display: block; }
#panel {
position: absolute; top: 14px; left: 14px;
background: rgba(13,17,23,0.92); border: 1px solid #30363d;
border-radius: 10px; padding: 16px 20px; z-index: 10;
width: 340px; backdrop-filter: blur(8px);
}
#panel h2 { font-size: 14px; color: #58a6ff; margin-bottom: 10px; letter-spacing: 0.5px; }
.row { display: flex; align-items: center; gap: 8px; margin: 6px 0; font-size: 12px; }
.row label { width: 70px; color: #8b949e; flex-shrink: 0; }
.row .val { color: #f0f6fc; font-variant-numeric: tabular-nums; }
#transport {
margin-top: 12px; display: flex; align-items: center; gap: 8px;
}
#transport button {
background: #21262d; color: #c9d1d9; border: 1px solid #30363d;
padding: 6px 14px; border-radius: 6px; cursor: pointer;
font-family: inherit; font-size: 12px; transition: background 0.15s;
}
#transport button:hover { background: #30363d; }
#transport button.active { background: #1f6feb; border-color: #1f6feb; color: #fff; }
#scrubber {
width: 100%; margin-top: 8px;
-webkit-appearance: none; appearance: none;
height: 6px; border-radius: 3px; background: #21262d; outline: none;
}
#scrubber::-webkit-slider-thumb {
-webkit-appearance: none; width: 14px; height: 14px;
border-radius: 50%; background: #58a6ff; cursor: pointer;
}
#speed-ctrl { margin-top: 6px; }
#speed-ctrl select {
background: #21262d; color: #c9d1d9; border: 1px solid #30363d;
padding: 4px 8px; border-radius: 4px; font-family: inherit; font-size: 11px;
}
#frame-counter {
font-size: 11px; color: #8b949e; margin-top: 6px;
font-variant-numeric: tabular-nums;
}
.legend { display: flex; align-items: center; gap: 6px; margin: 3px 0; font-size: 11px; }
.dot { width: 10px; height: 10px; border-radius: 50%; display: inline-block; }
</style>
<link href="https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;600&display=swap" rel="stylesheet">
</head>
<body>
<div id="panel">
<h2>DATASET REPLAY — EE FRAME</h2>
<div style="font-size:11px;color:#8b949e;margin-bottom:8px;">glannuzel/grabette-dataset · episode 0</div>
<div class="legend"><span class="dot" style="background:#ff6b6b"></span> EE target (dataset)</div>
<div class="legend"><span class="dot" style="background:#ffd43b"></span> Trajectory (past)</div>
<div class="legend"><span class="dot" style="background:#30363d"></span> Trajectory (future)</div>
<div class="row"><label>x</label><span class="val" id="v-x">—</span></div>
<div class="row"><label>y</label><span class="val" id="v-y">—</span></div>
<div class="row"><label>z</label><span class="val" id="v-z">—</span></div>
<div class="row"><label>ax</label><span class="val" id="v-ax">—</span></div>
<div class="row"><label>ay</label><span class="val" id="v-ay">—</span></div>
<div class="row"><label>az</label><span class="val" id="v-az">—</span></div>
<div class="row"><label>gripper</label><span class="val" id="v-grip">—</span></div>
<div id="transport">
<button id="btn-play" onclick="togglePlay()">▶ Play</button>
<button onclick="stepFrame(-1)">◀</button>
<button onclick="stepFrame(1)">▶</button>
<button onclick="resetPlay()">⟳</button>
</div>
<input type="range" id="scrubber" min="0" max="1" value="0" step="1" />
<div id="speed-ctrl">
<label style="font-size:11px;color:#8b949e;">Speed:</label>
<select id="speed-select" onchange="setSpeed(this.value)">
<option value="0.25">0.25×</option>
<option value="0.5">0.5×</option>
<option value="1" selected>1×</option>
<option value="2">2×</option>
<option value="4">4×</option>
</select>
</div>
<div id="frame-counter">Frame 0 / 0 · 0.00s</div>
</div>
<script type="importmap">
{
"imports": {
"three": "https://cdn.jsdelivr.net/npm/three@0.169.0/build/three.module.js",
"three/examples/jsm/": "https://cdn.jsdelivr.net/npm/three@0.169.0/examples/jsm/"
}
}
</script>
<script type="module">
import * as THREE from 'three';
import { OrbitControls } from 'three/examples/jsm/controls/OrbitControls.js';
import { STLLoader } from 'three/examples/jsm/loaders/STLLoader.js';
let trajectory = null;
let currentFrame = 0;
let playing = false;
let speed = 1.0;
let lastTime = 0;
let accumulator = 0;
// Anchor: EE tip world position at zero-joint pose (in Y-up Three.js space)
const eeAnchor = new THREE.Vector3();
// Z-up → Y-up rotation (same as robotGroup): -90° around X
const zUpToYUp = new THREE.Quaternion().setFromAxisAngle(new THREE.Vector3(1, 0, 0), -Math.PI / 2);
const scene = new THREE.Scene();
scene.background = new THREE.Color(0x0d1117);
const camera = new THREE.PerspectiveCamera(50, window.innerWidth / window.innerHeight, 0.01, 100);
const renderer = new THREE.WebGLRenderer({ antialias: true });
renderer.setSize(window.innerWidth, window.innerHeight);
renderer.setPixelRatio(window.devicePixelRatio);
renderer.shadowMap.enabled = true;
document.body.appendChild(renderer.domElement);
const controls = new OrbitControls(camera, renderer.domElement);
controls.enableDamping = true;
controls.dampingFactor = 0.08;
scene.add(new THREE.AmbientLight(0xffffff, 0.8));
const dirLight = new THREE.DirectionalLight(0xffffff, 1.4);
dirLight.position.set(2, 4, 3);
scene.add(dirLight);
scene.add(new THREE.DirectionalLight(0x8899cc, 0.6).translateX(-2).translateY(1).translateZ(-3));
scene.add(new THREE.DirectionalLight(0xffffff, 0.5).translateY(-1).translateZ(2));
const grid = new THREE.GridHelper(2, 20, 0x21262d, 0x161b22);
scene.add(grid);
scene.add(new THREE.AxesHelper(0.15));
// EE marker
const eeMarker = new THREE.Mesh(
new THREE.SphereGeometry(0.012, 20, 20),
new THREE.MeshStandardMaterial({ color: 0xff6b6b, emissive: 0xff6b6b, emissiveIntensity: 0.7 })
);
scene.add(eeMarker);
eeMarker.add(new THREE.AxesHelper(0.06));
// Trajectory lines
const MAX_POINTS = 2000;
const pastGeo = new THREE.BufferGeometry();
pastGeo.setAttribute('position', new THREE.Float32BufferAttribute(new Float32Array(MAX_POINTS * 3), 3));
const pastLine = new THREE.Line(pastGeo, new THREE.LineBasicMaterial({ color: 0xffd43b, linewidth: 2 }));
scene.add(pastLine);
const futureGeo = new THREE.BufferGeometry();
futureGeo.setAttribute('position', new THREE.Float32BufferAttribute(new Float32Array(MAX_POINTS * 3), 3));
const futureLine = new THREE.Line(futureGeo, new THREE.LineBasicMaterial({ color: 0x30363d, linewidth: 1 }));
scene.add(futureLine);
// URDF
const stlLoader = new STLLoader();
const robotGroup = new THREE.Group();
// URDF is Z-up; Three.js is Y-up → rotate -90° around X
robotGroup.rotation.x = -Math.PI / 2;
scene.add(robotGroup);
let urdfLinks = {};
function rotvecToQuat(ax, ay, az) {
const angle = Math.sqrt(ax * ax + ay * ay + az * az);
if (angle < 1e-8) return new THREE.Quaternion();
return new THREE.Quaternion().setFromAxisAngle(
new THREE.Vector3(ax / angle, ay / angle, az / angle), angle
);
}
async function loadURDF() {
const resp = await fetch('./openarm_bimanual_pybullet.urdf');
const text = await resp.text();
const xml = new DOMParser().parseFromString(text, 'text/xml');
const links = {};
for (const linkEl of xml.querySelectorAll('link')) {
const name = linkEl.getAttribute('name');
const group = new THREE.Group();
group.name = name;
const visual = linkEl.querySelector('visual');
if (visual) {
const meshEl = visual.querySelector('mesh');
const originEl = visual.querySelector('origin');
if (meshEl) {
const filename = meshEl.getAttribute('filename');
const scaleStr = meshEl.getAttribute('scale');
const sc = scaleStr ? scaleStr.split(' ').map(Number) : [1, 1, 1];
let xyz = [0, 0, 0];
if (originEl && originEl.getAttribute('xyz'))
xyz = originEl.getAttribute('xyz').split(' ').map(Number);
if (filename.endsWith('.stl')) {
try {
const geo = await new Promise((res, rej) =>
stlLoader.load(filename, res, undefined, rej));
const mesh = new THREE.Mesh(geo, new THREE.MeshStandardMaterial({
color: 0x8899bb, metalness: 0.3, roughness: 0.5,
}));
mesh.scale.set(sc[0], sc[1], sc[2]);
mesh.position.set(xyz[0], xyz[1], xyz[2]);
group.add(mesh);
} catch (e) { /* skip missing mesh */ }
}
}
}
links[name] = group;
}
const rootLinks = new Set(Object.keys(links));
for (const jointEl of xml.querySelectorAll('joint')) {
const parentName = jointEl.querySelector('parent').getAttribute('link');
const childName = jointEl.querySelector('child').getAttribute('link');
rootLinks.delete(childName);
const originEl = jointEl.querySelector('origin');
let xyz = [0, 0, 0], rpy = [0, 0, 0];
if (originEl) {
if (originEl.getAttribute('xyz')) xyz = originEl.getAttribute('xyz').split(' ').map(Number);
if (originEl.getAttribute('rpy')) rpy = originEl.getAttribute('rpy').split(' ').map(Number);
}
const parent = links[parentName];
const child = links[childName];
if (!parent || !child) continue;
child.position.set(xyz[0], xyz[1], xyz[2]);
if (rpy[0] || rpy[1] || rpy[2])
child.rotation.set(rpy[0], rpy[1], rpy[2], 'XYZ');
parent.add(child);
}
for (const n of rootLinks)
if (links[n]) robotGroup.add(links[n]);
// EE target marker on the URDF
const eeTargetLink = links['openarm_right_ee_target'];
if (eeTargetLink) {
eeTargetLink.add(new THREE.Mesh(
new THREE.TorusGeometry(0.02, 0.002, 8, 32),
new THREE.MeshStandardMaterial({ color: 0xffaa00, emissive: 0xffaa00, emissiveIntensity: 0.5 })
));
eeTargetLink.add(new THREE.AxesHelper(0.05));
}
urdfLinks = links;
}
async function loadTrajectory() {
const resp = await fetch('./trajectory_ep0.json');
trajectory = await resp.json();
document.getElementById('scrubber').max = trajectory.num_frames - 1;
document.getElementById('scrubber').value = 0;
}
function computeOffset() {
if (!trajectory || !urdfLinks['openarm_right_ee_target']) return;
robotGroup.updateMatrixWorld(true);
const eeLink = urdfLinks['openarm_right_ee_target'];
eeLink.getWorldPosition(eeAnchor);
controls.target.copy(eeAnchor);
camera.position.set(eeAnchor.x + 0.8, eeAnchor.y + 0.3, eeAnchor.z + 0.0);
controls.update();
updateFrame(0);
}
function mapFramePos(f) {
const f0 = trajectory.frames[0];
const delta = new THREE.Vector3(f.x - f0.x, f.y - f0.y, f.z - f0.z);
delta.applyQuaternion(zUpToYUp);
return delta.add(eeAnchor);
}
function updateFrame(idx) {
if (!trajectory) return;
currentFrame = Math.max(0, Math.min(idx, trajectory.num_frames - 1));
const f = trajectory.frames[currentFrame];
const pos = mapFramePos(f);
eeMarker.position.copy(pos);
// Orientation: rotate the dataset axis-angle into Y-up space
const q = rotvecToQuat(f.ax, f.ay, f.az);
eeMarker.quaternion.copy(zUpToYUp).multiply(q);
// Past trajectory
const pastArr = pastGeo.attributes.position.array;
let pi = 0;
for (let i = 0; i <= currentFrame && i < MAX_POINTS; i++) {
const p = mapFramePos(trajectory.frames[i]);
pastArr[pi++] = p.x; pastArr[pi++] = p.y; pastArr[pi++] = p.z;
}
pastGeo.setDrawRange(0, Math.min(currentFrame + 1, MAX_POINTS));
pastGeo.attributes.position.needsUpdate = true;
// Future trajectory
const futArr = futureGeo.attributes.position.array;
let fi = 0;
for (let i = currentFrame; i < trajectory.num_frames && (i - currentFrame) < MAX_POINTS; i++) {
const p = mapFramePos(trajectory.frames[i]);
futArr[fi++] = p.x; futArr[fi++] = p.y; futArr[fi++] = p.z;
}
futureGeo.setDrawRange(0, Math.min(trajectory.num_frames - currentFrame, MAX_POINTS));
futureGeo.attributes.position.needsUpdate = true;
// UI
document.getElementById('v-x').textContent = pos.x.toFixed(4);
document.getElementById('v-y').textContent = pos.y.toFixed(4);
document.getElementById('v-z').textContent = pos.z.toFixed(4);
document.getElementById('v-ax').textContent = f.ax.toFixed(4);
document.getElementById('v-ay').textContent = f.ay.toFixed(4);
document.getElementById('v-az').textContent = f.az.toFixed(4);
document.getElementById('v-grip').textContent =
`p=${f.proximal.toFixed(2)} d=${f.distal.toFixed(2)}`;
document.getElementById('scrubber').value = currentFrame;
const timeS = (currentFrame / trajectory.fps).toFixed(2);
document.getElementById('frame-counter').textContent =
`Frame ${currentFrame} / ${trajectory.num_frames - 1} · ${timeS}s`;
}
// Playback controls
window.togglePlay = function() {
playing = !playing;
const btn = document.getElementById('btn-play');
btn.textContent = playing ? '⏸ Pause' : '▶ Play';
btn.classList.toggle('active', playing);
if (playing) { lastTime = performance.now(); accumulator = 0; }
};
window.stepFrame = function(delta) {
playing = false;
document.getElementById('btn-play').textContent = '▶ Play';
document.getElementById('btn-play').classList.remove('active');
updateFrame(currentFrame + delta);
};
window.resetPlay = function() {
playing = false;
document.getElementById('btn-play').textContent = '▶ Play';
document.getElementById('btn-play').classList.remove('active');
updateFrame(0);
};
window.setSpeed = function(v) { speed = parseFloat(v); };
document.getElementById('scrubber').addEventListener('input', (e) => {
updateFrame(parseInt(e.target.value));
});
window.addEventListener('resize', () => {
camera.aspect = window.innerWidth / window.innerHeight;
camera.updateProjectionMatrix();
renderer.setSize(window.innerWidth, window.innerHeight);
});
function animate(now) {
requestAnimationFrame(animate);
controls.update();
if (playing && trajectory) {
const dt = (now - lastTime) / 1000;
lastTime = now;
accumulator += dt * speed;
const frameDuration = 1.0 / trajectory.fps;
while (accumulator >= frameDuration) {
accumulator -= frameDuration;
if (currentFrame < trajectory.num_frames - 1) {
updateFrame(currentFrame + 1);
} else {
playing = false;
document.getElementById('btn-play').textContent = '▶ Play';
document.getElementById('btn-play').classList.remove('active');
break;
}
}
}
renderer.render(scene, camera);
}
requestAnimationFrame(animate);
Promise.all([loadURDF(), loadTrajectory()])
.then(() => computeOffset())
.catch(err => console.error(err));
</script>
</body>
</html>

311
src/lerobot/robots/openarm_follower/urdf/viewer.html Normal file
View File

@@ -0,0 +1,311 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8" />
<title>OpenArm URDF Viewer</title>
<style>
* { margin: 0; padding: 0; box-sizing: border-box; }
body { background: #1a1a2e; overflow: hidden; font-family: 'IBM Plex Mono', monospace; }
canvas { display: block; }
#info {
position: absolute; top: 16px; left: 16px;
color: #e0e0e0; font-size: 13px; line-height: 1.6;
background: rgba(0,0,0,0.7); padding: 14px 18px; border-radius: 8px;
border: 1px solid #333; max-width: 340px; z-index: 10;
}
#info h2 { font-size: 15px; color: #fff; margin-bottom: 8px; }
.legend { display: flex; align-items: center; gap: 8px; margin: 4px 0; }
.dot { width: 12px; height: 12px; border-radius: 50%; display: inline-block; flex-shrink: 0; }
.dot-red { background: #ff4444; }
.dot-green { background: #44ff44; }
.dot-blue { background: #4488ff; }
#frame-select { margin-top: 10px; }
#frame-select button {
background: #333; color: #e0e0e0; border: 1px solid #555;
padding: 6px 10px; margin: 2px; border-radius: 4px; cursor: pointer;
font-family: inherit; font-size: 12px;
}
#frame-select button:hover { background: #555; }
#frame-select button.active { background: #4488ff; color: #fff; border-color: #4488ff; }
#status { margin-top: 8px; font-size: 11px; color: #888; }
</style>
<link href="https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;600&display=swap" rel="stylesheet">
</head>
<body>
<div id="info">
<h2>OpenArm Right Arm — EE Frame Options</h2>
<div class="legend"><span class="dot dot-red"></span> openarm_right_link7 (wrist output)</div>
<div class="legend"><span class="dot" style="background:#ffaa00"></span> openarm_right_ee_target (+5cm)</div>
<div class="legend"><span class="dot dot-green"></span> openarm_right_hand (+10cm)</div>
<div class="legend"><span class="dot dot-blue"></span> openarm_right_hand_tcp (+18cm)</div>
<div id="frame-select">
<button onclick="focusFrame('link7')" class="active">link7</button>
<button onclick="focusFrame('ee_target')">ee_target</button>
<button onclick="focusFrame('hand')">hand</button>
<button onclick="focusFrame('tcp')">hand_tcp</button>
</div>
<div id="status">Loading URDF...</div>
<p style="margin-top:8px;font-size:11px;color:#888;">Drag to orbit · Scroll to zoom · Right-drag to pan</p>
</div>
<script type="importmap">
{
"imports": {
"three": "https://cdn.jsdelivr.net/npm/three@0.169.0/build/three.module.js",
"three/examples/jsm/": "https://cdn.jsdelivr.net/npm/three@0.169.0/examples/jsm/"
}
}
</script>
<script type="module">
import * as THREE from 'three';
import { OrbitControls } from 'three/examples/jsm/controls/OrbitControls.js';
import { STLLoader } from 'three/examples/jsm/loaders/STLLoader.js';
const statusEl = document.getElementById('status');
const scene = new THREE.Scene();
scene.background = new THREE.Color(0x1a1a2e);
const camera = new THREE.PerspectiveCamera(50, window.innerWidth / window.innerHeight, 0.01, 100);
camera.position.set(0.8, 1.0, 1.8);
const renderer = new THREE.WebGLRenderer({ antialias: true });
renderer.setSize(window.innerWidth, window.innerHeight);
renderer.setPixelRatio(window.devicePixelRatio);
renderer.shadowMap.enabled = true;
document.body.appendChild(renderer.domElement);
const controls = new OrbitControls(camera, renderer.domElement);
controls.target.set(0, 0, 0.9);
controls.enableDamping = true;
controls.dampingFactor = 0.08;
controls.update();
// Lighting
scene.add(new THREE.AmbientLight(0xffffff, 0.6));
const dirLight = new THREE.DirectionalLight(0xffffff, 1.2);
dirLight.position.set(3, 5, 4);
scene.add(dirLight);
scene.add(new THREE.DirectionalLight(0x8888ff, 0.4).translateX(-2).translateY(1).translateZ(-3));
// Ground grid
scene.add(new THREE.GridHelper(4, 40, 0x333355, 0x222244));
scene.add(new THREE.AxesHelper(0.3));
// Parse URDF manually — build the kinematic tree and load STL meshes
const stlLoader = new STLLoader();
const robotGroup = new THREE.Group();
scene.add(robotGroup);
async function loadURDF() {
const resp = await fetch('./openarm_bimanual_pybullet.urdf');
const text = await resp.text();
const parser = new DOMParser();
const xml = parser.parseFromString(text, 'text/xml');
// Parse links and joints
const links = {};
const joints = [];
for (const linkEl of xml.querySelectorAll('link')) {
const name = linkEl.getAttribute('name');
const group = new THREE.Group();
group.name = name;
// Try to load visual mesh
const visual = linkEl.querySelector('visual');
if (visual) {
const meshEl = visual.querySelector('mesh');
const originEl = visual.querySelector('origin');
if (meshEl) {
const filename = meshEl.getAttribute('filename');
const scaleStr = meshEl.getAttribute('scale');
const scale = scaleStr ? scaleStr.split(' ').map(Number) : [1, 1, 1];
let xyz = [0, 0, 0];
if (originEl && originEl.getAttribute('xyz')) {
xyz = originEl.getAttribute('xyz').split(' ').map(Number);
}
if (filename.endsWith('.stl')) {
try {
const geo = await new Promise((resolve, reject) => {
stlLoader.load(filename, resolve, undefined, reject);
});
const mat = new THREE.MeshStandardMaterial({
color: 0x6688aa,
metalness: 0.3,
roughness: 0.6,
transparent: true,
opacity: 0.7,
});
const mesh = new THREE.Mesh(geo, mat);
mesh.scale.set(scale[0], scale[1], scale[2]);
mesh.position.set(xyz[0], xyz[1], xyz[2]);
group.add(mesh);
} catch (e) {
// Mesh file not found, skip
}
}
}
}
links[name] = group;
}
// Parse joints and build hierarchy
for (const jointEl of xml.querySelectorAll('joint')) {
const name = jointEl.getAttribute('name');
const type = jointEl.getAttribute('type');
const parentName = jointEl.querySelector('parent').getAttribute('link');
const childName = jointEl.querySelector('child').getAttribute('link');
const originEl = jointEl.querySelector('origin');
let xyz = [0, 0, 0];
let rpy = [0, 0, 0];
if (originEl) {
if (originEl.getAttribute('xyz')) xyz = originEl.getAttribute('xyz').split(' ').map(Number);
if (originEl.getAttribute('rpy')) rpy = originEl.getAttribute('rpy').split(' ').map(Number);
}
joints.push({ name, type, parentName, childName, xyz, rpy });
}
// Build tree
const rootLinks = new Set(Object.keys(links));
for (const j of joints) {
rootLinks.delete(j.childName);
}
for (const j of joints) {
const parent = links[j.parentName];
const child = links[j.childName];
if (parent && child) {
child.position.set(j.xyz[0], j.xyz[1], j.xyz[2]);
if (j.rpy[0] || j.rpy[1] || j.rpy[2]) {
child.rotation.set(j.rpy[0], j.rpy[1], j.rpy[2], 'XYZ');
}
parent.add(child);
}
}
// Add root links to scene
for (const name of rootLinks) {
if (links[name]) robotGroup.add(links[name]);
}
// Place markers at the three EE frame candidates
const targets = {
link7: 'openarm_right_link7',
ee_target: 'openarm_right_ee_target',
hand: 'openarm_right_hand',
tcp: 'openarm_right_hand_tcp',
};
const colors = { link7: 0xff4444, ee_target: 0xffaa00, hand: 0x44ff44, tcp: 0x4488ff };
const labels = { link7: 'link7', ee_target: 'ee_target', hand: 'hand', tcp: 'hand_tcp' };
for (const [key, linkName] of Object.entries(targets)) {
const link = links[linkName];
if (!link) continue;
// Marker sphere
const sphere = new THREE.Mesh(
new THREE.SphereGeometry(0.018, 16, 16),
new THREE.MeshStandardMaterial({ color: colors[key], emissive: colors[key], emissiveIntensity: 0.6 })
);
link.add(sphere);
// Ring around sphere for visibility
const ring = new THREE.Mesh(
new THREE.TorusGeometry(0.03, 0.003, 8, 32),
new THREE.MeshStandardMaterial({ color: colors[key], emissive: colors[key], emissiveIntensity: 0.4 })
);
link.add(ring);
// Axes helper
link.add(new THREE.AxesHelper(0.08));
// Sprite label
const canvas = document.createElement('canvas');
canvas.width = 512; canvas.height = 80;
const ctx = canvas.getContext('2d');
ctx.font = 'bold 36px IBM Plex Mono, monospace';
ctx.fillStyle = '#' + colors[key].toString(16).padStart(6, '0');
ctx.fillText(labels[key], 4, 50);
const tex = new THREE.CanvasTexture(canvas);
const sprite = new THREE.Sprite(new THREE.SpriteMaterial({ map: tex, depthTest: false }));
sprite.scale.set(0.3, 0.05, 1);
sprite.position.set(0.06, 0.0, 0.03);
link.add(sprite);
}
// Dashed lines between markers (in world space)
robotGroup.updateMatrixWorld(true);
const positions = {};
for (const [key, linkName] of Object.entries(targets)) {
const link = links[linkName];
if (link) {
const wp = new THREE.Vector3();
link.getWorldPosition(wp);
positions[key] = wp;
}
}
function addDashedLine(from, to) {
const geo = new THREE.BufferGeometry().setFromPoints([from, to]);
const mat = new THREE.LineDashedMaterial({ color: 0xaaaaaa, dashSize: 0.012, gapSize: 0.008 });
const line = new THREE.Line(geo, mat);
line.computeLineDistances();
scene.add(line);
}
if (positions.link7 && positions.hand) addDashedLine(positions.link7, positions.hand);
if (positions.hand && positions.tcp) addDashedLine(positions.hand, positions.tcp);
// Store for focus buttons
window._framePositions = positions;
window._links = links;
window._targets = targets;
// Focus on the hand area
if (positions.hand) {
controls.target.copy(positions.hand);
camera.position.set(positions.hand.x + 0.5, positions.hand.y + 0.4, positions.hand.z + 0.5);
controls.update();
}
const meshCount = robotGroup.children.length;
statusEl.textContent = `Loaded. Right arm chain visible with ${Object.keys(links).length} links.`;
}
window.focusFrame = function(key) {
const pos = window._framePositions?.[key];
if (!pos) return;
controls.target.copy(pos);
camera.position.set(pos.x + 0.35, pos.y + 0.25, pos.z + 0.35);
controls.update();
document.querySelectorAll('#frame-select button').forEach(b => b.classList.remove('active'));
event.target.classList.add('active');
};
window.addEventListener('resize', () => {
camera.aspect = window.innerWidth / window.innerHeight;
camera.updateProjectionMatrix();
renderer.setSize(window.innerWidth, window.innerHeight);
});
function animate() {
requestAnimationFrame(animate);
controls.update();
renderer.render(scene, camera);
}
animate();
loadURDF().catch(err => {
statusEl.textContent = `Error: ${err.message}`;
console.error(err);
});
</script>
</body>
</html>

180
src/lerobot/robots/so_follower/robot_kinematic_processor.py
View File

@@ -255,19 +255,16 @@ class InverseKinematicsEEToJoints(RobotActionProcessorStep):
"""
Computes desired joint positions from a target end-effector pose using inverse kinematics (IK).
This step translates a Cartesian command (position and orientation of the end-effector) into
the corresponding joint-space commands for each motor.
Attributes:
kinematics: The robot's kinematic model for inverse kinematics.
motor_names: A list of motor names for which to compute joint positions.
q_curr: Internal state storing the last joint positions, used as an initial guess for the IK solver.
motor_names: Arm joint names for IK computation.
gripper_names: Gripper joint name(s). ee.gripper_pos is written to all of them.
initial_guess_current_joints: If True, use the robot's current joint state as the IK guess.
If False, use the solution from the previous step.
"""
kinematics: RobotKinematics
motor_names: list[str]
gripper_names: list[str] = field(default_factory=lambda: ["gripper"])
q_curr: np.ndarray | None = field(default=None, init=False, repr=False)
initial_guess_current_joints: bool = True
@@ -278,63 +275,73 @@ class InverseKinematicsEEToJoints(RobotActionProcessorStep):
wx = action.pop("ee.wx")
wy = action.pop("ee.wy")
wz = action.pop("ee.wz")
gripper_pos = action.pop("ee.gripper_pos")
if None in (x, y, z, wx, wy, wz, gripper_pos):
raise ValueError(
"Missing required end-effector pose components: ee.x, ee.y, ee.z, ee.wx, ee.wy, ee.wz, ee.gripper_pos must all be present in action"
)
ee_keys = [x, y, z, wx, wy, wz]
if self.gripper_names:
gripper_pos = action.pop("ee.gripper_pos")
ee_keys.append(gripper_pos)
if None in ee_keys:
raise ValueError("Missing required end-effector pose components in action")
observation = self.transition.get(TransitionKey.OBSERVATION).copy()
if observation is None:
raise ValueError("Joints observation is require for computing robot kinematics")
raise ValueError("Joints observation is required for computing robot kinematics")
q_raw = np.array(
[float(v) for k, v in observation.items() if isinstance(k, str) and k.endswith(".pos")],
[
float(v)
for k, v in observation.items()
if isinstance(k, str) and k.endswith(".pos") and k.removesuffix(".pos") in self.motor_names
],
dtype=float,
)
if q_raw is None:
raise ValueError("Joints observation is require for computing robot kinematics")
if self.initial_guess_current_joints: # Use current joints as initial guess
if self.initial_guess_current_joints:
self.q_curr = q_raw
else: # Use previous ik solution as initial guess
else:
if self.q_curr is None:
self.q_curr = q_raw
# Build desired 4x4 transform from pos + rotvec (twist)
t_des = np.eye(4, dtype=float)
t_des[:3, :3] = Rotation.from_rotvec([wx, wy, wz]).as_matrix()
t_des[:3, 3] = [x, y, z]
# Compute inverse kinematics
q_target = self.kinematics.inverse_kinematics(self.q_curr, t_des)
self.q_curr = q_target
# TODO: This is sentitive to order of motor_names = q_target mapping
for i, name in enumerate(self.motor_names):
if name != "gripper":
action[f"{name}.pos"] = float(q_target[i])
else:
action["gripper.pos"] = float(gripper_pos)
action[f"{name}.pos"] = float(q_target[i])
if self.gripper_names:
for gname in self.gripper_names:
action[f"{gname}.pos"] = float(gripper_pos)
# When gripper_names is empty, gripper keys (e.g. proximal.pos, distal.pos)
# are already in the action dict as absolute positions — left untouched.
return action
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
for feat in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
ee_feats = ["x", "y", "z", "wx", "wy", "wz"]
if self.gripper_names:
ee_feats.append("gripper_pos")
for feat in ee_feats:
features[PipelineFeatureType.ACTION].pop(f"ee.{feat}", None)
for name in self.motor_names:
features[PipelineFeatureType.ACTION][f"{name}.pos"] = PolicyFeature(
type=FeatureType.ACTION, shape=(1,)
)
for name in self.gripper_names:
features[PipelineFeatureType.ACTION][f"{name}.pos"] = PolicyFeature(
type=FeatureType.ACTION, shape=(1,)
)
return features
def reset(self):
"""Resets the initial guess for the IK solver."""
self.q_curr = None
@@ -402,24 +409,39 @@ class GripperVelocityToJoint(RobotActionProcessorStep):
def compute_forward_kinematics_joints_to_ee(
joints: dict[str, Any], kinematics: RobotKinematics, motor_names: list[str]
joints: dict[str, Any],
kinematics: RobotKinematics,
motor_names: list[str],
gripper_names: list[str] | None = None,
) -> dict[str, Any]:
if gripper_names is None:
gripper_names = ["gripper"]
motor_joint_values = [joints[f"{n}.pos"] for n in motor_names]
q = np.array(motor_joint_values, dtype=float)
t = kinematics.forward_kinematics(q)
pos = t[:3, 3]
tw = Rotation.from_matrix(t[:3, :3]).as_rotvec()
gripper_pos = joints["gripper.pos"]
for n in motor_names:
joints.pop(f"{n}.pos")
joints["ee.x"] = float(pos[0])
joints["ee.y"] = float(pos[1])
joints["ee.z"] = float(pos[2])
joints["ee.wx"] = float(tw[0])
joints["ee.wy"] = float(tw[1])
joints["ee.wz"] = float(tw[2])
joints["ee.gripper_pos"] = float(gripper_pos)
# When gripper_names is non-empty, fold them into ee.gripper_pos (e.g. SO100).
# When empty, gripper joints pass through as-is (absolute position control).
if gripper_names:
gripper_pos = joints[f"{gripper_names[0]}.pos"]
for n in gripper_names:
joints.pop(f"{n}.pos", None)
joints["ee.gripper_pos"] = float(gripper_pos)
return joints
@@ -429,27 +451,33 @@ class ForwardKinematicsJointsToEEObservation(ObservationProcessorStep):
"""
Computes the end-effector pose from joint positions using forward kinematics (FK).
This step is typically used to add the robot's Cartesian pose to the observation space,
which can be useful for visualization or as an input to a policy.
Attributes:
kinematics: The robot's kinematic model.
motor_names: Arm joint names used for FK computation.
gripper_names: Gripper joint name(s) to fold into ee.gripper_pos.
Empty list means gripper joints pass through as absolute positions.
"""
kinematics: RobotKinematics
motor_names: list[str]
gripper_names: list[str] = field(default_factory=lambda: ["gripper"])
def observation(self, observation: RobotObservation) -> RobotObservation:
return compute_forward_kinematics_joints_to_ee(observation, self.kinematics, self.motor_names)
return compute_forward_kinematics_joints_to_ee(
observation, self.kinematics, self.motor_names, self.gripper_names
)
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
# We only use the ee pose in the dataset, so we don't need the joint positions
for n in self.motor_names:
features[PipelineFeatureType.OBSERVATION].pop(f"{n}.pos", None)
# We specify the dataset features of this step that we want to be stored in the dataset
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
ee_keys = ["x", "y", "z", "wx", "wy", "wz"]
if self.gripper_names:
for n in self.gripper_names:
features[PipelineFeatureType.OBSERVATION].pop(f"{n}.pos", None)
ee_keys.append("gripper_pos")
for k in ee_keys:
features[PipelineFeatureType.OBSERVATION][f"ee.{k}"] = PolicyFeature(
type=FeatureType.STATE, shape=(1,)
)
@@ -462,27 +490,33 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
"""
Computes the end-effector pose from joint positions using forward kinematics (FK).
This step is typically used to add the robot's Cartesian pose to the observation space,
which can be useful for visualization or as an input to a policy.
Attributes:
kinematics: The robot's kinematic model.
motor_names: Arm joint names used for FK computation.
gripper_names: Gripper joint name(s) to fold into ee.gripper_pos.
Empty list means gripper joints pass through as absolute positions.
"""
kinematics: RobotKinematics
motor_names: list[str]
gripper_names: list[str] = field(default_factory=lambda: ["gripper"])
def action(self, action: RobotAction) -> RobotAction:
return compute_forward_kinematics_joints_to_ee(action, self.kinematics, self.motor_names)
return compute_forward_kinematics_joints_to_ee(
action, self.kinematics, self.motor_names, self.gripper_names
)
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
# We only use the ee pose in the dataset, so we don't need the joint positions
for n in self.motor_names:
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
# We specify the dataset features of this step that we want to be stored in the dataset
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
ee_keys = ["x", "y", "z", "wx", "wy", "wz"]
if self.gripper_names:
for n in self.gripper_names:
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
ee_keys.append("gripper_pos")
for k in ee_keys:
features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
type=FeatureType.STATE, shape=(1,)
)
@@ -494,13 +528,14 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
class ForwardKinematicsJointsToEE(ProcessorStep):
kinematics: RobotKinematics
motor_names: list[str]
gripper_names: list[str] = field(default_factory=lambda: ["gripper"])
def __post_init__(self):
self.joints_to_ee_action_processor = ForwardKinematicsJointsToEEAction(
kinematics=self.kinematics, motor_names=self.motor_names
kinematics=self.kinematics, motor_names=self.motor_names, gripper_names=self.gripper_names
)
self.joints_to_ee_observation_processor = ForwardKinematicsJointsToEEObservation(
kinematics=self.kinematics, motor_names=self.motor_names
kinematics=self.kinematics, motor_names=self.motor_names, gripper_names=self.gripper_names
)
def __call__(self, transition: EnvTransition) -> EnvTransition:
@@ -524,13 +559,13 @@ class ForwardKinematicsJointsToEE(ProcessorStep):
@dataclass
class InverseKinematicsRLStep(ProcessorStep):
"""
Computes desired joint positions from a target end-effector pose using inverse kinematics (IK).
This is modified from the InverseKinematicsEEToJoints step to be used in the RL pipeline.
IK step for the RL pipeline. Same logic as InverseKinematicsEEToJoints but
operates on EnvTransition directly and stores the IK solution.
"""
kinematics: RobotKinematics
motor_names: list[str]
gripper_names: list[str] = field(default_factory=lambda: ["gripper"])
q_curr: np.ndarray | None = field(default=None, init=False, repr=False)
initial_guess_current_joints: bool = True
@@ -538,7 +573,7 @@ class InverseKinematicsRLStep(ProcessorStep):
new_transition = dict(transition)
action = new_transition.get(TransitionKey.ACTION)
if action is None:
raise ValueError("Action is required for InverseKinematicsEEToJoints")
raise ValueError("Action is required for InverseKinematicsRLStep")
action = dict(action)
x = action.pop("ee.x")
@@ -547,45 +582,46 @@ class InverseKinematicsRLStep(ProcessorStep):
wx = action.pop("ee.wx")
wy = action.pop("ee.wy")
wz = action.pop("ee.wz")
gripper_pos = action.pop("ee.gripper_pos")
if None in (x, y, z, wx, wy, wz, gripper_pos):
raise ValueError(
"Missing required end-effector pose components: ee.x, ee.y, ee.z, ee.wx, ee.wy, ee.wz, ee.gripper_pos must all be present in action"
)
ee_keys = [x, y, z, wx, wy, wz]
if self.gripper_names:
gripper_pos = action.pop("ee.gripper_pos")
ee_keys.append(gripper_pos)
if None in ee_keys:
raise ValueError("Missing required end-effector pose components in action")
observation = new_transition.get(TransitionKey.OBSERVATION).copy()
if observation is None:
raise ValueError("Joints observation is require for computing robot kinematics")
raise ValueError("Joints observation is required for computing robot kinematics")
q_raw = np.array(
[float(v) for k, v in observation.items() if isinstance(k, str) and k.endswith(".pos")],
[
float(v)
for k, v in observation.items()
if isinstance(k, str) and k.endswith(".pos") and k.removesuffix(".pos") in self.motor_names
],
dtype=float,
)
if q_raw is None:
raise ValueError("Joints observation is require for computing robot kinematics")
if self.initial_guess_current_joints: # Use current joints as initial guess
if self.initial_guess_current_joints:
self.q_curr = q_raw
else: # Use previous ik solution as initial guess
else:
if self.q_curr is None:
self.q_curr = q_raw
# Build desired 4x4 transform from pos + rotvec (twist)
t_des = np.eye(4, dtype=float)
t_des[:3, :3] = Rotation.from_rotvec([wx, wy, wz]).as_matrix()
t_des[:3, 3] = [x, y, z]
# Compute inverse kinematics
q_target = self.kinematics.inverse_kinematics(self.q_curr, t_des)
self.q_curr = q_target
# TODO: This is sentitive to order of motor_names = q_target mapping
for i, name in enumerate(self.motor_names):
if name != "gripper":
action[f"{name}.pos"] = float(q_target[i])
else:
action["gripper.pos"] = float(gripper_pos)
action[f"{name}.pos"] = float(q_target[i])
if self.gripper_names:
for gname in self.gripper_names:
action[f"{gname}.pos"] = float(gripper_pos)
new_transition[TransitionKey.ACTION] = action
complementary_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
@@ -596,16 +632,22 @@ class InverseKinematicsRLStep(ProcessorStep):
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
for feat in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
ee_feats = ["x", "y", "z", "wx", "wy", "wz"]
if self.gripper_names:
ee_feats.append("gripper_pos")
for feat in ee_feats:
features[PipelineFeatureType.ACTION].pop(f"ee.{feat}", None)
for name in self.motor_names:
features[PipelineFeatureType.ACTION][f"{name}.pos"] = PolicyFeature(
type=FeatureType.ACTION, shape=(1,)
)
for name in self.gripper_names:
features[PipelineFeatureType.ACTION][f"{name}.pos"] = PolicyFeature(
type=FeatureType.ACTION, shape=(1,)
)
return features
def reset(self):
"""Resets the initial guess for the IK solver."""
self.q_curr = None