mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-16 09:09:48 +00:00
Add relative code
This commit is contained in:
Pepijn
@@ -0,0 +1,257 @@
|
||||
#!/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.
|
||||
|
||||
"""
|
||||
Convert a joint-space OpenArms dataset to end-effector space.
|
||||
|
||||
For each frame, converts joint positions to EE poses (x, y, z, wx, wy, wz) using FK.
|
||||
Grippers are kept as-is. Applies to both observation.state and action.
|
||||
|
||||
Example usage:
|
||||
python examples/openarms/convert_joints_to_ee.py \
|
||||
--input-dataset lerobot-data-collection/rac_blackf0 \
|
||||
--output-repo-id my_user/rac_blackf0_ee \
|
||||
--output-dir ./outputs/rac_blackf0_ee
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from tqdm import tqdm
|
||||
|
||||
from lerobot.datasets.compute_stats import get_feature_stats
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
|
||||
from lerobot.datasets.utils import write_info, write_stats
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.utils.rotation import Rotation
|
||||
|
||||
DEFAULT_URDF = "src/lerobot/robots/openarms/urdf/openarm_bimanual_pybullet.urdf"
|
||||
DEFAULT_LEFT_EE_FRAME = "openarm_left_hand_tcp"
|
||||
DEFAULT_RIGHT_EE_FRAME = "openarm_right_hand_tcp"
|
||||
|
||||
LEFT_URDF_JOINTS = [f"openarm_left_joint{i}" for i in range(1, 8)]
|
||||
RIGHT_URDF_JOINTS = [f"openarm_right_joint{i}" for i in range(1, 8)]
|
||||
|
||||
JOINT_NAMES = [f"joint_{i}" for i in range(1, 8)]
|
||||
EE_COMPONENTS = ["x", "y", "z", "wx", "wy", "wz"]
|
||||
|
||||
|
||||
def compute_fk_for_arm(kinematics: RobotKinematics, joint_values: np.ndarray) -> dict[str, float]:
|
||||
"""Compute FK for one arm, returns EE pose as dict."""
|
||||
t = kinematics.forward_kinematics(joint_values)
|
||||
pos = t[:3, 3]
|
||||
rotvec = Rotation.from_matrix(t[:3, :3]).as_rotvec()
|
||||
return {
|
||||
"x": float(pos[0]),
|
||||
"y": float(pos[1]),
|
||||
"z": float(pos[2]),
|
||||
"wx": float(rotvec[0]),
|
||||
"wy": float(rotvec[1]),
|
||||
"wz": float(rotvec[2]),
|
||||
}
|
||||
|
||||
|
||||
def convert_joints_to_ee(
|
||||
values: np.ndarray,
|
||||
names: list[str],
|
||||
left_kin: RobotKinematics,
|
||||
right_kin: RobotKinematics,
|
||||
) -> tuple[np.ndarray, list[str]]:
|
||||
"""
|
||||
Convert joint values to EE values.
|
||||
|
||||
Args:
|
||||
values: Array of shape (N,) with joint values for one frame
|
||||
names: List of feature names corresponding to values
|
||||
left_kin: Left arm kinematics solver
|
||||
right_kin: Right arm kinematics solver
|
||||
|
||||
Returns:
|
||||
(new_values, new_names) with joints replaced by EE poses
|
||||
"""
|
||||
name_to_idx = {n: i for i, n in enumerate(names)}
|
||||
|
||||
new_values = []
|
||||
new_names = []
|
||||
|
||||
for prefix, kinematics in [("right", right_kin), ("left", left_kin)]:
|
||||
joint_vals = []
|
||||
for jname in JOINT_NAMES:
|
||||
key = f"{prefix}_{jname}.pos"
|
||||
if key in name_to_idx:
|
||||
joint_vals.append(values[name_to_idx[key]])
|
||||
|
||||
if len(joint_vals) == 7:
|
||||
ee_pose = compute_fk_for_arm(kinematics, np.array(joint_vals, dtype=float))
|
||||
for comp in EE_COMPONENTS:
|
||||
new_names.append(f"{prefix}_ee.{comp}")
|
||||
new_values.append(ee_pose[comp])
|
||||
|
||||
gripper_key = f"{prefix}_gripper.pos"
|
||||
if gripper_key in name_to_idx:
|
||||
new_names.append(f"{prefix}_ee.gripper_pos")
|
||||
new_values.append(values[name_to_idx[gripper_key]])
|
||||
|
||||
return np.array(new_values, dtype=np.float32), new_names
|
||||
|
||||
|
||||
def transform_feature_info(old_info: dict, new_names: list[str]) -> dict:
|
||||
"""Create new feature info with EE names instead of joint names."""
|
||||
return {
|
||||
"dtype": old_info.get("dtype", "float32"),
|
||||
"shape": (len(new_names),),
|
||||
"names": new_names,
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Convert joint-space dataset to EE-space")
|
||||
parser.add_argument("--input-dataset", type=str, required=True, help="Input dataset repo ID")
|
||||
parser.add_argument("--output-repo-id", type=str, required=True, help="Output dataset repo ID")
|
||||
parser.add_argument("--output-dir", type=str, required=True, help="Output directory")
|
||||
parser.add_argument("--urdf", type=str, default=DEFAULT_URDF, help="Path to URDF file")
|
||||
parser.add_argument("--left-ee-frame", type=str, default=DEFAULT_LEFT_EE_FRAME)
|
||||
parser.add_argument("--right-ee-frame", type=str, default=DEFAULT_RIGHT_EE_FRAME)
|
||||
parser.add_argument("--push-to-hub", action="store_true", help="Push converted dataset to HF Hub")
|
||||
args = parser.parse_args()
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
if output_dir.exists():
|
||||
shutil.rmtree(output_dir)
|
||||
|
||||
urdf_path = args.urdf
|
||||
if not Path(urdf_path).is_absolute():
|
||||
urdf_path = str(Path(__file__).parent.parent.parent / urdf_path)
|
||||
|
||||
print(f"Loading dataset: {args.input_dataset}")
|
||||
dataset = LeRobotDataset(args.input_dataset)
|
||||
|
||||
print(f"Initializing kinematics from {urdf_path}")
|
||||
left_kin = RobotKinematics(urdf_path, args.left_ee_frame, LEFT_URDF_JOINTS)
|
||||
right_kin = RobotKinematics(urdf_path, args.right_ee_frame, RIGHT_URDF_JOINTS)
|
||||
|
||||
action_info = dataset.meta.features.get("action", {})
|
||||
state_info = dataset.meta.features.get("observation.state", {})
|
||||
action_names = action_info.get("names", [])
|
||||
state_names = state_info.get("names", [])
|
||||
|
||||
print(f"Original action names ({len(action_names)}): {action_names[:8]}...")
|
||||
print(f"Original state names ({len(state_names)}): {state_names[:8]}...")
|
||||
|
||||
sample_action = np.zeros(len(action_names), dtype=np.float32)
|
||||
_, new_action_names = convert_joints_to_ee(sample_action, action_names, left_kin, right_kin)
|
||||
sample_state = np.zeros(len(state_names), dtype=np.float32)
|
||||
_, new_state_names = convert_joints_to_ee(sample_state, state_names, left_kin, right_kin)
|
||||
|
||||
print(f"New action names ({len(new_action_names)}): {new_action_names}")
|
||||
print(f"New state names ({len(new_state_names)}): {new_state_names}")
|
||||
|
||||
new_features = dataset.meta.features.copy()
|
||||
new_features["action"] = transform_feature_info(action_info, new_action_names)
|
||||
new_features["observation.state"] = transform_feature_info(state_info, new_state_names)
|
||||
|
||||
new_meta = LeRobotDatasetMetadata.create(
|
||||
repo_id=args.output_repo_id,
|
||||
fps=dataset.meta.fps,
|
||||
features=new_features,
|
||||
robot_type=dataset.meta.robot_type,
|
||||
root=output_dir,
|
||||
use_videos=len(dataset.meta.video_keys) > 0,
|
||||
)
|
||||
|
||||
data_dir = dataset.root / "data"
|
||||
parquet_files = sorted(data_dir.glob("*/*.parquet"))
|
||||
print(f"Processing {len(parquet_files)} parquet files...")
|
||||
|
||||
all_actions = []
|
||||
all_states = []
|
||||
|
||||
for src_path in tqdm(parquet_files, desc="Converting"):
|
||||
df = pd.read_parquet(src_path).reset_index(drop=True)
|
||||
|
||||
new_actions = []
|
||||
new_states = []
|
||||
|
||||
for idx in range(len(df)):
|
||||
action_vals = np.array(df.iloc[idx]["action"], dtype=np.float32)
|
||||
state_vals = np.array(df.iloc[idx]["observation.state"], dtype=np.float32)
|
||||
|
||||
new_action, _ = convert_joints_to_ee(action_vals, action_names, left_kin, right_kin)
|
||||
new_state, _ = convert_joints_to_ee(state_vals, state_names, left_kin, right_kin)
|
||||
|
||||
new_actions.append(new_action.tolist())
|
||||
new_states.append(new_state.tolist())
|
||||
all_actions.append(new_action)
|
||||
all_states.append(new_state)
|
||||
|
||||
df["action"] = new_actions
|
||||
df["observation.state"] = new_states
|
||||
|
||||
relative_path = src_path.relative_to(dataset.root)
|
||||
out_path = output_dir / relative_path
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
df.to_parquet(out_path)
|
||||
|
||||
print("Computing statistics...")
|
||||
all_actions_arr = np.stack(all_actions)
|
||||
all_states_arr = np.stack(all_states)
|
||||
|
||||
stats = {}
|
||||
stats["action"] = get_feature_stats(all_actions_arr, axis=0, keepdims=True)
|
||||
stats["observation.state"] = get_feature_stats(all_states_arr, axis=0, keepdims=True)
|
||||
write_stats(stats, output_dir)
|
||||
|
||||
print("Updating metadata...")
|
||||
new_meta.info["total_episodes"] = dataset.meta.total_episodes
|
||||
new_meta.info["total_frames"] = dataset.meta.total_frames
|
||||
new_meta.info["total_tasks"] = dataset.meta.total_tasks
|
||||
write_info(new_meta.info, output_dir)
|
||||
|
||||
print("Copying episode metadata...")
|
||||
src_episodes_dir = dataset.root / "meta" / "episodes"
|
||||
dst_episodes_dir = output_dir / "meta" / "episodes"
|
||||
if src_episodes_dir.exists():
|
||||
shutil.copytree(src_episodes_dir, dst_episodes_dir, dirs_exist_ok=True)
|
||||
|
||||
print("Copying tasks metadata...")
|
||||
src_tasks = dataset.root / "meta" / "tasks.parquet"
|
||||
dst_tasks = output_dir / "meta" / "tasks.parquet"
|
||||
if src_tasks.exists():
|
||||
shutil.copy2(src_tasks, dst_tasks)
|
||||
|
||||
if dataset.meta.video_keys:
|
||||
print("Copying videos...")
|
||||
src_videos = dataset.root / "videos"
|
||||
dst_videos = output_dir / "videos"
|
||||
if src_videos.exists():
|
||||
shutil.copytree(src_videos, dst_videos, dirs_exist_ok=True)
|
||||
|
||||
print(f"\nDone! Dataset saved to: {output_dir}")
|
||||
print(f"Repo ID: {args.output_repo_id}")
|
||||
|
||||
if args.push_to_hub:
|
||||
print("\nPushing to Hub...")
|
||||
output_dataset = LeRobotDataset(args.output_repo_id, root=output_dir)
|
||||
output_dataset.push_to_hub()
|
||||
print(f"Pushed to: https://huggingface.co/datasets/{args.output_repo_id}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -0,0 +1,403 @@
|
||||
#!/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.
|
||||
|
||||
"""
|
||||
OpenArms End-Effector Replay Example with Visualization
|
||||
|
||||
Replays a dataset recorded with absolute joint positions by:
|
||||
1. Converting joint positions to EE poses using FK
|
||||
2. Converting EE poses back to joint positions using IK
|
||||
3. Sending joint commands to the robot OR visualizing in simulation
|
||||
|
||||
Supports three modes:
|
||||
- real: Send commands to physical robot
|
||||
- sim: Visualize in simulation only (no robot required)
|
||||
- both: Real robot + visualization
|
||||
|
||||
Example usage:
|
||||
python examples/openarms/replay_ee.py --mode sim
|
||||
python examples/openarms/replay_ee.py --mode real
|
||||
python examples/openarms/replay_ee.py --mode both --visualizer meshcat
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from os.path import dirname, expanduser
|
||||
|
||||
import numpy as np
|
||||
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import RobotAction, RobotObservation, RobotProcessorPipeline
|
||||
from lerobot.processor.converters import (
|
||||
robot_action_observation_to_transition,
|
||||
robot_action_to_transition,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
from lerobot.robots.openarms.robot_kinematic_processor import (
|
||||
BimanualEEBoundsAndSafety,
|
||||
BimanualForwardKinematicsJointsToEE,
|
||||
BimanualInverseKinematicsEEToJoints,
|
||||
)
|
||||
from lerobot.utils.constants import ACTION
|
||||
from lerobot.utils.robot_utils import precise_sleep
|
||||
|
||||
|
||||
# Default configuration
|
||||
DEFAULT_EPISODE_IDX = 0
|
||||
DEFAULT_DATASET = "lerobot-data-collection/rac_blackf0"
|
||||
DEFAULT_URDF = "src/lerobot/robots/openarms/urdf/openarm_bimanual_pybullet.urdf"
|
||||
DEFAULT_LEFT_EE_FRAME = "openarm_left_hand_tcp"
|
||||
DEFAULT_RIGHT_EE_FRAME = "openarm_right_hand_tcp"
|
||||
|
||||
# Motor names as used in the dataset actions (e.g., left_joint_1.pos)
|
||||
MOTOR_NAMES = ["joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6", "joint_7", "gripper"]
|
||||
|
||||
# URDF joint names (no underscore between "joint" and number)
|
||||
LEFT_URDF_JOINTS = [f"openarm_left_joint{i}" for i in range(1, 8)]
|
||||
RIGHT_URDF_JOINTS = [f"openarm_right_joint{i}" for i in range(1, 8)]
|
||||
|
||||
|
||||
class MeshcatVisualizer:
|
||||
"""Lightweight URDF visualizer using pinocchio + meshcat."""
|
||||
|
||||
def __init__(self, urdf_path: str):
|
||||
import pinocchio as pin
|
||||
from pinocchio.visualize import MeshcatVisualizer as PinMeshcat
|
||||
|
||||
urdf_dir = dirname(urdf_path)
|
||||
self.model, self.collision_model, self.visual_model = pin.buildModelsFromUrdf(
|
||||
urdf_path, urdf_dir, pin.JointModelFreeFlyer()
|
||||
)
|
||||
self.data = self.model.createData()
|
||||
|
||||
self.viz = PinMeshcat(self.model, self.collision_model, self.visual_model)
|
||||
self.viz.initViewer(open=True)
|
||||
self.viz.loadViewerModel()
|
||||
|
||||
# Build joint name mapping: dataset name -> pinocchio joint index
|
||||
# Dataset uses: left_joint_1, right_joint_2, etc.
|
||||
# URDF uses: openarm_left_joint1, openarm_right_joint2, etc.
|
||||
self.joint_map = {}
|
||||
for jid in range(1, self.model.njoints):
|
||||
urdf_name = self.model.names[jid] # e.g., "openarm_left_joint1"
|
||||
# Extract side and number
|
||||
if "left_joint" in urdf_name:
|
||||
num = urdf_name.split("joint")[-1] # "1"
|
||||
dataset_name = f"left_joint_{num}"
|
||||
self.joint_map[dataset_name] = jid
|
||||
elif "right_joint" in urdf_name:
|
||||
num = urdf_name.split("joint")[-1]
|
||||
dataset_name = f"right_joint_{num}"
|
||||
self.joint_map[dataset_name] = jid
|
||||
|
||||
print(f" Meshcat viewer opened (mapped {len(self.joint_map)} joints)")
|
||||
print(f" Joint map: {list(self.joint_map.keys())[:4]}...")
|
||||
print(" Waiting for meshcat to load...")
|
||||
time.sleep(3) # Give meshcat time to load meshes
|
||||
self._first_update = True
|
||||
|
||||
def update(self, joint_positions: dict[str, float]):
|
||||
"""Update visualization with new joint positions."""
|
||||
if self._first_update:
|
||||
pos_keys = [k for k in joint_positions.keys() if k.endswith(".pos")]
|
||||
print(f" First update keys: {pos_keys[:4]}...")
|
||||
# Print sample values
|
||||
for k in pos_keys[:2]:
|
||||
print(f" {k} = {joint_positions[k]:.2f}")
|
||||
|
||||
# Build configuration vector (base pose + joints)
|
||||
# Free flyer base: [x, y, z, qx, qy, qz, qw]
|
||||
q = np.zeros(self.model.nq)
|
||||
q[3:7] = [0, 0, 0, 1] # Identity quaternion
|
||||
|
||||
matched = 0
|
||||
# Map joint positions using pre-built mapping
|
||||
for name, pos in joint_positions.items():
|
||||
if not name.endswith(".pos"):
|
||||
continue
|
||||
joint_name = name.removesuffix(".pos") # e.g., "left_joint_1"
|
||||
|
||||
jid = self.joint_map.get(joint_name)
|
||||
if jid is not None:
|
||||
idx = self.model.idx_qs[jid]
|
||||
if idx < len(q):
|
||||
q[idx] = np.deg2rad(pos)
|
||||
matched += 1
|
||||
|
||||
if self._first_update:
|
||||
print(f" Matched {matched} joints, q[7:14] = {q[7:14]}")
|
||||
self._first_update = False
|
||||
|
||||
self.viz.display(q)
|
||||
|
||||
|
||||
class RerunVisualizer:
|
||||
"""Rerun-based visualizer for plots and EE trajectories."""
|
||||
|
||||
def __init__(self, urdf_path: str = None, session_name: str = "openarms_replay"):
|
||||
import rerun as rr
|
||||
self.rr = rr
|
||||
rr.init(session_name)
|
||||
rr.spawn(memory_limit="10%")
|
||||
print(" Rerun viewer spawned (plots only, use --visualizer meshcat for 3D robot)")
|
||||
|
||||
def update(self, joint_positions: dict[str, float], ee_poses: dict[str, float], frame_idx: int):
|
||||
"""Log joint positions and EE poses."""
|
||||
self.rr.set_time("frame", sequence=frame_idx)
|
||||
|
||||
# Log EE positions as colored spheres
|
||||
for prefix, color in [("left", [255, 100, 100]), ("right", [100, 100, 255])]:
|
||||
x, y, z = ee_poses.get(f"{prefix}_ee.x"), ee_poses.get(f"{prefix}_ee.y"), ee_poses.get(f"{prefix}_ee.z")
|
||||
if None not in (x, y, z):
|
||||
self.rr.log(f"ee/{prefix}", self.rr.Points3D([[x, y, z]], colors=[color], radii=[0.02]))
|
||||
|
||||
# Log joint positions as time series
|
||||
for name, pos in joint_positions.items():
|
||||
if name.endswith(".pos"):
|
||||
self.rr.log(f"joints/{name}", self.rr.Scalars(pos))
|
||||
|
||||
# Log EE poses as time series
|
||||
for name, val in ee_poses.items():
|
||||
self.rr.log(f"ee_plots/{name}", self.rr.Scalars(val))
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="OpenArms EE Replay with Visualization")
|
||||
parser.add_argument("--mode", choices=["real", "sim", "both"], default="sim",
|
||||
help="Execution mode: real (robot), sim (visualization), both")
|
||||
parser.add_argument("--visualizer", choices=["meshcat", "rerun", "none"], default="meshcat",
|
||||
help="Visualization backend (meshcat shows 3D robot, rerun shows plots)")
|
||||
parser.add_argument("--dataset", type=str, default=DEFAULT_DATASET,
|
||||
help="Dataset repo ID")
|
||||
parser.add_argument("--episode", type=int, default=DEFAULT_EPISODE_IDX,
|
||||
help="Episode index to replay")
|
||||
parser.add_argument("--urdf", type=str, default=DEFAULT_URDF,
|
||||
help="Path to URDF file")
|
||||
parser.add_argument("--left-ee-frame", type=str, default=DEFAULT_LEFT_EE_FRAME,
|
||||
help="Left arm end-effector frame name in URDF")
|
||||
parser.add_argument("--right-ee-frame", type=str, default=DEFAULT_RIGHT_EE_FRAME,
|
||||
help="Right arm end-effector frame name in URDF")
|
||||
parser.add_argument("--port-left", type=str, default="can2",
|
||||
help="CAN port for left arm")
|
||||
parser.add_argument("--port-right", type=str, default="can3",
|
||||
help="CAN port for right arm")
|
||||
parser.add_argument("--speed", type=float, default=1.0,
|
||||
help="Playback speed multiplier")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
use_robot = args.mode in ["real", "both"]
|
||||
use_viz = args.mode in ["sim", "both"] and args.visualizer != "none"
|
||||
|
||||
print("=" * 70)
|
||||
print("OpenArms EE Replay (FK -> IK Pipeline)")
|
||||
print("=" * 70)
|
||||
print(f"\nMode: {args.mode}")
|
||||
print(f"Visualizer: {args.visualizer}")
|
||||
print(f"Dataset: {args.dataset}")
|
||||
print(f"Episode: {args.episode}")
|
||||
print(f"Speed: {args.speed}x")
|
||||
print("=" * 70)
|
||||
|
||||
robot = None
|
||||
viz = None
|
||||
|
||||
# Resolve URDF path (handle relative and ~ paths)
|
||||
from pathlib import Path
|
||||
urdf_path = args.urdf
|
||||
if urdf_path.startswith("~"):
|
||||
urdf_path = expanduser(urdf_path)
|
||||
elif not Path(urdf_path).is_absolute():
|
||||
# Relative to workspace root
|
||||
urdf_path = str(Path(__file__).parent.parent.parent / urdf_path)
|
||||
|
||||
# Initialize robot if needed
|
||||
if use_robot:
|
||||
from lerobot.robots.openarms.config_openarms_follower import OpenArmsFollowerConfig
|
||||
from lerobot.robots.openarms.openarms_follower import OpenArmsFollower
|
||||
|
||||
print("\n[1/5] Initializing robot...")
|
||||
robot_config = OpenArmsFollowerConfig(
|
||||
port_left=args.port_left,
|
||||
port_right=args.port_right,
|
||||
can_interface="socketcan",
|
||||
id="openarms_follower",
|
||||
disable_torque_on_disconnect=True,
|
||||
max_relative_target=10.0,
|
||||
)
|
||||
robot = OpenArmsFollower(robot_config)
|
||||
else:
|
||||
print("\n[1/5] Skipping robot (sim mode)")
|
||||
|
||||
# Initialize visualizer if needed
|
||||
if use_viz:
|
||||
print(f"\n[2/5] Initializing {args.visualizer} visualizer...")
|
||||
if args.visualizer == "meshcat":
|
||||
viz = MeshcatVisualizer(urdf_path)
|
||||
elif args.visualizer == "rerun":
|
||||
viz = RerunVisualizer(urdf_path)
|
||||
else:
|
||||
print("\n[2/5] Skipping visualization")
|
||||
|
||||
# Initialize kinematics with URDF joint names
|
||||
print("\n[3/5] Initializing kinematics solvers...")
|
||||
|
||||
left_kinematics = RobotKinematics(
|
||||
urdf_path=urdf_path,
|
||||
target_frame_name=args.left_ee_frame,
|
||||
joint_names=LEFT_URDF_JOINTS,
|
||||
)
|
||||
right_kinematics = RobotKinematics(
|
||||
urdf_path=urdf_path,
|
||||
target_frame_name=args.right_ee_frame,
|
||||
joint_names=RIGHT_URDF_JOINTS,
|
||||
)
|
||||
|
||||
# Build pipelines - use motor names without gripper for the processor
|
||||
motor_names_no_gripper = [n for n in MOTOR_NAMES if n != "gripper"]
|
||||
|
||||
joints_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
|
||||
steps=[
|
||||
BimanualForwardKinematicsJointsToEE(
|
||||
left_kinematics=left_kinematics,
|
||||
right_kinematics=right_kinematics,
|
||||
motor_names=MOTOR_NAMES,
|
||||
),
|
||||
],
|
||||
to_transition=robot_action_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
|
||||
ee_to_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
steps=[
|
||||
BimanualEEBoundsAndSafety(
|
||||
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
|
||||
max_ee_step_m=0.10,
|
||||
),
|
||||
BimanualInverseKinematicsEEToJoints(
|
||||
left_kinematics=left_kinematics,
|
||||
right_kinematics=right_kinematics,
|
||||
motor_names=MOTOR_NAMES,
|
||||
initial_guess_current_joints=False,
|
||||
),
|
||||
],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
|
||||
# Load dataset
|
||||
print(f"\n[4/5] Loading dataset '{args.dataset}'...")
|
||||
dataset = LeRobotDataset(args.dataset, episodes=[args.episode])
|
||||
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == args.episode)
|
||||
|
||||
if len(episode_frames) == 0:
|
||||
raise ValueError(f"No frames found for episode {args.episode}")
|
||||
|
||||
print(f" Found {len(episode_frames)} frames at {dataset.fps} FPS")
|
||||
|
||||
action_features = dataset.features.get(ACTION, {})
|
||||
action_names = action_features.get("names", [])
|
||||
actions = episode_frames.select_columns(ACTION)
|
||||
|
||||
# Connect robot if needed
|
||||
if use_robot:
|
||||
print("\n[5/5] Connecting to robot...")
|
||||
robot.connect(calibrate=False)
|
||||
if not robot.is_connected:
|
||||
raise RuntimeError("Robot failed to connect!")
|
||||
else:
|
||||
print("\n[5/5] Skipping robot connection (sim mode)")
|
||||
|
||||
print("\n" + "=" * 70)
|
||||
print(f"Ready to replay! Mode: {args.mode}")
|
||||
print("=" * 70)
|
||||
|
||||
if use_robot:
|
||||
input("\nPress ENTER to start...")
|
||||
else:
|
||||
print("\nStarting visualization playback...")
|
||||
time.sleep(1)
|
||||
|
||||
# Simulated observation for sim-only mode
|
||||
sim_obs = {f"{prefix}_{motor}.pos": 0.0
|
||||
for prefix in ["left", "right"]
|
||||
for motor in MOTOR_NAMES}
|
||||
|
||||
try:
|
||||
for idx in range(len(episode_frames)):
|
||||
loop_start = time.perf_counter()
|
||||
|
||||
# Get observation
|
||||
if use_robot:
|
||||
robot_obs = robot.get_observation()
|
||||
else:
|
||||
robot_obs = sim_obs.copy()
|
||||
|
||||
# Build joint action from dataset
|
||||
action_array = actions[idx][ACTION]
|
||||
joint_action = {}
|
||||
for i, name in enumerate(action_names):
|
||||
if name.endswith(".pos"):
|
||||
joint_action[name] = float(action_array[i])
|
||||
|
||||
# Convert: joints -> EE (FK)
|
||||
ee_action = joints_to_ee(joint_action.copy())
|
||||
|
||||
# Convert: EE -> joints (IK)
|
||||
final_joint_action = ee_to_joints((ee_action.copy(), robot_obs))
|
||||
|
||||
# Update simulated observation for next iteration
|
||||
if not use_robot:
|
||||
sim_obs.update(final_joint_action)
|
||||
|
||||
# Send to robot
|
||||
if use_robot:
|
||||
robot.send_action(final_joint_action)
|
||||
|
||||
# Update visualization with ORIGINAL dataset trajectory
|
||||
if viz:
|
||||
if isinstance(viz, MeshcatVisualizer):
|
||||
viz.update(joint_action) # Use original, not FK->IK reconstructed
|
||||
elif isinstance(viz, RerunVisualizer):
|
||||
viz.update(joint_action, ee_action, idx)
|
||||
|
||||
# Maintain replay rate
|
||||
loop_duration = time.perf_counter() - loop_start
|
||||
dt_s = (1.0 / dataset.fps / args.speed) - loop_duration
|
||||
if dt_s > 0:
|
||||
precise_sleep(dt_s)
|
||||
|
||||
if (idx + 1) % 100 == 0:
|
||||
progress = (idx + 1) / len(episode_frames) * 100
|
||||
print(f"Progress: {idx + 1}/{len(episode_frames)} ({progress:.1f}%)")
|
||||
|
||||
print(f"\n✓ Replayed {len(episode_frames)} frames")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n\nReplay interrupted")
|
||||
finally:
|
||||
if use_robot and robot:
|
||||
print("\nDisconnecting robot...")
|
||||
robot.disconnect()
|
||||
print("✓ Done!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user