| from __future__ import annotations |
|
|
| import json |
| import os |
| import shlex |
| import shutil |
| import subprocess |
| import sys |
| import time |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import mujoco |
| import numpy as np |
| import torch |
| from lerobot.configs.policies import PreTrainedConfig |
| from lerobot.configs.types import FeatureType |
| from lerobot.datasets.lerobot_dataset import LeRobotDataset |
| from lerobot.policies.factory import get_policy_class, make_pre_post_processors |
|
|
| try: |
| import mujoco.viewer |
| except Exception: |
| pass |
|
|
|
|
| JOINT_NAMES = [ |
| "shoulder_pan.pos", |
| "shoulder_lift.pos", |
| "elbow_flex.pos", |
| "wrist_flex.pos", |
| "wrist_roll.pos", |
| "gripper.pos", |
| ] |
|
|
| HOME_STATE = np.array([0.0, -35.0, 75.0, -40.0, 0.0, 70.0], dtype=np.float64) |
| TABLE_Z = 0.025 |
| ATTACH_OFFSET = np.array([0.0, 0.0, -0.055], dtype=np.float64) |
| GRIP_CLOSE_THRESHOLD = 25.0 |
| GRIP_OPEN_THRESHOLD = 40.0 |
| GRASP_DISTANCE_THRESHOLD = 0.045 |
| SUCCESS_XY_THRESHOLD = 0.05 |
|
|
|
|
| @dataclass |
| class DatasetConfig: |
| repo_id: str |
| root: Path |
| episodes: int |
| fps: int |
| width: int |
| height: int |
| image_key: str |
| task: str |
| robot_type: str |
| base_height: float |
| seed: int |
| overwrite: bool |
| show_viewer: bool |
|
|
|
|
| @dataclass |
| class VizConfig: |
| episode_index: int |
|
|
|
|
| @dataclass |
| class TrainConfig: |
| output_dir: Path |
| job_name: str |
| device: str |
| batch_size: int |
| num_workers: int |
| steps: int |
| save_freq: int |
| log_freq: int |
| eval_freq: int |
| seed: int |
| wandb_enable: bool |
| resume: bool |
|
|
|
|
| @dataclass |
| class EvalConfig: |
| policy_root: Path | None |
| episodes: int |
| max_steps: int |
| device: str |
| base_height: float |
| seed: int |
| show_viewer: bool |
| step_sleep: float |
| task: str |
| robot_type: str |
|
|
|
|
| @dataclass |
| class AppConfig: |
| config_path: Path |
| dataset: DatasetConfig |
| viz: VizConfig |
| train: TrainConfig |
| eval: EvalConfig |
|
|
|
|
| def resolve_path(base_dir: Path, value: str) -> Path: |
| path = Path(value).expanduser() |
| if path.is_absolute(): |
| return path |
| return (base_dir / path).resolve() |
|
|
|
|
| def load_config(config_path: str | Path | None = None) -> AppConfig: |
| path = Path(config_path or Path(__file__).with_name("config.json")).expanduser().resolve() |
| raw = json.loads(path.read_text()) |
| base_dir = path.parent |
|
|
| dataset = DatasetConfig( |
| repo_id=raw["dataset"]["repo_id"], |
| root=resolve_path(base_dir, raw["dataset"]["root"]), |
| episodes=int(raw["dataset"]["episodes"]), |
| fps=int(raw["dataset"]["fps"]), |
| width=int(raw["dataset"]["width"]), |
| height=int(raw["dataset"]["height"]), |
| image_key=raw["dataset"]["image_key"], |
| task=raw["dataset"]["task"], |
| robot_type=raw["dataset"]["robot_type"], |
| base_height=float(raw["dataset"]["base_height"]), |
| seed=int(raw["dataset"]["seed"]), |
| overwrite=bool(raw["dataset"]["overwrite"]), |
| show_viewer=bool(raw["dataset"]["show_viewer"]), |
| ) |
| viz = VizConfig(episode_index=int(raw["viz"]["episode_index"])) |
| train = TrainConfig( |
| output_dir=resolve_path(base_dir, raw["train"]["output_dir"]), |
| job_name=raw["train"]["job_name"], |
| device=raw["train"]["device"], |
| batch_size=int(raw["train"]["batch_size"]), |
| num_workers=int(raw["train"]["num_workers"]), |
| steps=int(raw["train"]["steps"]), |
| save_freq=int(raw["train"]["save_freq"]), |
| log_freq=int(raw["train"]["log_freq"]), |
| eval_freq=int(raw["train"]["eval_freq"]), |
| seed=int(raw["train"]["seed"]), |
| wandb_enable=bool(raw["train"]["wandb_enable"]), |
| resume=bool(raw["train"]["resume"]), |
| ) |
| policy_root_value = raw["eval"]["policy_root"] |
| eval_cfg = EvalConfig( |
| policy_root=resolve_path(base_dir, policy_root_value) if policy_root_value else None, |
| episodes=int(raw["eval"]["episodes"]), |
| max_steps=int(raw["eval"]["max_steps"]), |
| device=raw["eval"]["device"], |
| base_height=float(raw["eval"]["base_height"]), |
| seed=int(raw["eval"]["seed"]), |
| show_viewer=bool(raw["eval"]["show_viewer"]), |
| step_sleep=float(raw["eval"]["step_sleep"]), |
| task=raw["eval"]["task"], |
| robot_type=raw["eval"]["robot_type"], |
| ) |
| return AppConfig(config_path=path, dataset=dataset, viz=viz, train=train, eval=eval_cfg) |
|
|
|
|
| def lerobot_env() -> dict[str, str]: |
| return os.environ.copy() |
|
|
|
|
| def maybe_reexec_with_mjpython(show_viewer: bool, env_key: str) -> None: |
| if not show_viewer or sys.platform != "darwin": |
| return |
| if os.environ.get(env_key) == "1": |
| return |
| if Path(sys.executable).name == "mjpython": |
| return |
|
|
| mjpython_path = Path(sys.executable).with_name("mjpython") |
| if not mjpython_path.exists(): |
| raise RuntimeError(f"`mjpython` not found next to {sys.executable}") |
|
|
| env = os.environ.copy() |
| env[env_key] = "1" |
| os.execve(str(mjpython_path), [str(mjpython_path), *sys.argv], env) |
|
|
|
|
| def bool_cli(value: bool) -> str: |
| return "true" if value else "false" |
|
|
|
|
| def parse_device(device: str) -> torch.device: |
| if device == "auto": |
| if torch.cuda.is_available(): |
| return torch.device("cuda") |
| if torch.backends.mps.is_available(): |
| return torch.device("mps") |
| return torch.device("cpu") |
| return torch.device(device) |
|
|
|
|
| def build_xml(base_height: float) -> str: |
| return f""" |
| <mujoco model="mujoco_pickplace_minimal"> |
| <compiler angle="radian"/> |
| <option timestep="0.01" gravity="0 0 -9.81"/> |
| <visual> |
| <headlight ambient="0.7 0.7 0.7" diffuse="0.7 0.7 0.7" specular="0.1 0.1 0.1"/> |
| <global offwidth="1024" offheight="1024"/> |
| </visual> |
| <asset> |
| <texture name="grid" type="2d" builtin="checker" rgb1="0.28 0.31 0.35" rgb2="0.18 0.20 0.24" width="256" height="256"/> |
| <material name="floor" texture="grid" texrepeat="6 6" reflectance="0.05"/> |
| <material name="arm" rgba="0.88 0.66 0.22 1"/> |
| <material name="metal" rgba="0.24 0.24 0.28 1"/> |
| <material name="finger" rgba="0.75 0.77 0.80 1"/> |
| <material name="cube" rgba="0.91 0.32 0.20 1"/> |
| </asset> |
| <worldbody> |
| <light pos="0 0 2.4" dir="0 0 -1"/> |
| <geom name="floor" type="plane" size="2 2 0.05" material="floor"/> |
| <site name="goal_site" pos="0.18 0.18 0.025" size="0.03 0.03 0.003" type="box" rgba="0.2 0.8 0.3 0.55"/> |
| <body name="cube" pos="0.22 -0.10 0.025"> |
| <freejoint name="cube_freejoint"/> |
| <geom type="box" size="0.025 0.025 0.025" material="cube" mass="0.05"/> |
| </body> |
| <body name="base" pos="0 0 {base_height:.3f}"> |
| <geom type="cylinder" size="0.09 0.05" material="metal"/> |
| <body name="shoulder_pan_link" pos="0 0 0.05"> |
| <joint name="shoulder_pan" type="hinge" axis="0 0 1" range="-3.14 3.14"/> |
| <geom type="capsule" fromto="0 0 0 0 0 0.10" size="0.035" material="arm"/> |
| <body name="shoulder_lift_link" pos="0 0 0.10"> |
| <joint name="shoulder_lift" type="hinge" axis="0 1 0" range="-2.7 2.7"/> |
| <geom type="capsule" fromto="0 0 0 0.18 0 0" size="0.03" material="arm"/> |
| <body name="elbow_flex_link" pos="0.18 0 0"> |
| <joint name="elbow_flex" type="hinge" axis="0 1 0" range="-2.7 2.7"/> |
| <geom type="capsule" fromto="0 0 0 0.16 0 0" size="0.026" material="arm"/> |
| <body name="wrist_flex_link" pos="0.16 0 0"> |
| <joint name="wrist_flex" type="hinge" axis="0 1 0" range="-2.7 2.7"/> |
| <geom type="capsule" fromto="0 0 0 0.10 0 0" size="0.021" material="arm"/> |
| <body name="wrist_roll_link" pos="0.10 0 0"> |
| <joint name="wrist_roll" type="hinge" axis="1 0 0" range="-3.14 3.14"/> |
| <geom type="capsule" fromto="0 0 0 0.06 0 0" size="0.017" material="arm"/> |
| <body name="tcp_body" pos="0.06 0 0"> |
| <geom type="box" size="0.02 0.014 0.014" material="metal"/> |
| <site name="tcp" pos="0.035 0 0" size="0.008" rgba="0.2 0.85 0.35 1"/> |
| <body name="left_finger" pos="0.012 0.015 0"> |
| <joint name="left_finger_slide" type="slide" axis="0 1 0" range="0 0.025"/> |
| <geom type="box" pos="0.025 0.012 0" size="0.025 0.004 0.006" material="finger"/> |
| </body> |
| <body name="right_finger" pos="0.012 -0.015 0"> |
| <joint name="right_finger_slide" type="slide" axis="0 -1 0" range="0 0.025"/> |
| <geom type="box" pos="0.025 -0.012 0" size="0.025 0.004 0.006" material="finger"/> |
| </body> |
| </body> |
| </body> |
| </body> |
| </body> |
| </body> |
| </body> |
| </body> |
| </worldbody> |
| </mujoco> |
| """ |
|
|
|
|
| def prepare_model(base_height: float) -> tuple[mujoco.MjModel, mujoco.MjData, dict[str, int]]: |
| model = mujoco.MjModel.from_xml_string(build_xml(base_height)) |
| data = mujoco.MjData(model) |
| ids = { |
| "tcp": mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "tcp"), |
| "goal_site": mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "goal_site"), |
| "cube_joint": mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "cube_freejoint"), |
| "shoulder_pan": mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "shoulder_pan"), |
| "shoulder_lift": mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "shoulder_lift"), |
| "elbow_flex": mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "elbow_flex"), |
| "wrist_flex": mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "wrist_flex"), |
| "wrist_roll": mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "wrist_roll"), |
| "left_finger": mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "left_finger_slide"), |
| "right_finger": mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "right_finger_slide"), |
| } |
| return model, data, ids |
|
|
|
|
| def set_arm_state(model: mujoco.MjModel, data: mujoco.MjData, ids: dict[str, int], joint_state: np.ndarray) -> None: |
| qpos = data.qpos |
| qpos[model.jnt_qposadr[ids["shoulder_pan"]]] = np.deg2rad(joint_state[0]) |
| qpos[model.jnt_qposadr[ids["shoulder_lift"]]] = np.deg2rad(joint_state[1]) |
| qpos[model.jnt_qposadr[ids["elbow_flex"]]] = np.deg2rad(joint_state[2]) |
| qpos[model.jnt_qposadr[ids["wrist_flex"]]] = np.deg2rad(joint_state[3]) |
| qpos[model.jnt_qposadr[ids["wrist_roll"]]] = np.deg2rad(joint_state[4]) |
| finger_slide = 0.003 + 0.022 * float(np.clip(joint_state[5] / 100.0, 0.0, 1.0)) |
| qpos[model.jnt_qposadr[ids["left_finger"]]] = finger_slide |
| qpos[model.jnt_qposadr[ids["right_finger"]]] = finger_slide |
|
|
|
|
| def set_cube_pose(model: mujoco.MjModel, data: mujoco.MjData, ids: dict[str, int], pos: np.ndarray) -> None: |
| adr = model.jnt_qposadr[ids["cube_joint"]] |
| data.qpos[adr : adr + 7] = np.array([pos[0], pos[1], pos[2], 1.0, 0.0, 0.0, 0.0], dtype=np.float64) |
|
|
|
|
| def render_frame(renderer: mujoco.Renderer, data: mujoco.MjData, base_height: float) -> np.ndarray: |
| camera = mujoco.MjvCamera() |
| camera.type = mujoco.mjtCamera.mjCAMERA_FREE |
| camera.lookat[:] = np.array([0.18, 0.0, base_height + 0.08], dtype=np.float64) |
| camera.distance = 1.0 |
| camera.azimuth = 140.0 |
| camera.elevation = -45.0 |
| renderer.update_scene(data, camera=camera) |
| return renderer.render() |
|
|
|
|
| def configure_viewer_camera(viewer, base_height: float) -> None: |
| viewer.cam.lookat[:] = np.array([0.18, 0.0, base_height + 0.10], dtype=np.float64) |
| viewer.cam.distance = 1.15 |
| viewer.cam.azimuth = 135.0 |
| viewer.cam.elevation = -32.0 |
|
|
|
|
| def sample_positions(rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]: |
| while True: |
| obj = np.array([rng.uniform(0.18, 0.28), rng.uniform(-0.12, 0.12), TABLE_Z], dtype=np.float64) |
| goal = np.array([rng.uniform(0.10, 0.25), rng.uniform(-0.18, 0.18), TABLE_Z], dtype=np.float64) |
| if np.linalg.norm(obj[:2] - goal[:2]) >= 0.14: |
| return obj, goal |
|
|
|
|
| def solve_ik(target_xyz: np.ndarray) -> np.ndarray: |
| shoulder_origin = np.array([0.0, 0.0, 0.50], dtype=np.float64) |
| link_1 = 0.18 |
| link_2 = 0.16 |
| tool = 0.195 |
| desired_pitch = -np.pi / 2.0 |
|
|
| rel = target_xyz - shoulder_origin |
| yaw = np.arctan2(rel[1], rel[0]) |
| radial = np.hypot(rel[0], rel[1]) |
| wrist_x = radial - tool * np.cos(desired_pitch) |
| wrist_z = rel[2] - tool * np.sin(desired_pitch) |
| dist_sq = wrist_x**2 + wrist_z**2 |
| cos_elbow = np.clip((dist_sq - link_1**2 - link_2**2) / (2.0 * link_1 * link_2), -1.0, 1.0) |
| elbow = np.arccos(cos_elbow) |
| shoulder = np.arctan2(wrist_z, wrist_x) - np.arctan2(link_2 * np.sin(elbow), link_1 + link_2 * np.cos(elbow)) |
| wrist = desired_pitch - shoulder - elbow |
| return np.rad2deg(np.array([yaw, -shoulder, -elbow, -wrist, 0.0], dtype=np.float64)) |
|
|
|
|
| def interpolate(start: np.ndarray, end: np.ndarray, steps: int) -> list[np.ndarray]: |
| if steps <= 1: |
| return [end.astype(np.float64)] |
| return [(start + (end - start) * (i / (steps - 1))).astype(np.float64) for i in range(steps)] |
|
|
|
|
| def make_episode_trajectory(obj: np.ndarray, goal: np.ndarray) -> list[tuple[np.ndarray, bool, np.ndarray]]: |
| open_grip = 70.0 |
| closed_grip = 8.0 |
|
|
| above_obj = np.concatenate([solve_ik(obj + np.array([0.0, 0.0, 0.12])), [open_grip]]) |
| grasp = np.concatenate([solve_ik(obj + np.array([0.0, 0.0, 0.045])), [open_grip]]) |
| close = grasp.copy() |
| close[-1] = closed_grip |
| lift = np.concatenate([solve_ik(obj + np.array([0.0, 0.0, 0.18])), [closed_grip]]) |
| above_goal = np.concatenate([solve_ik(goal + np.array([0.0, 0.0, 0.16])), [closed_grip]]) |
| place = np.concatenate([solve_ik(goal + np.array([0.0, 0.0, 0.05])), [closed_grip]]) |
| open_at_goal = place.copy() |
| open_at_goal[-1] = open_grip |
| retreat = np.concatenate([solve_ik(goal + np.array([0.0, 0.0, 0.18])), [open_grip]]) |
|
|
| segments = [ |
| (HOME_STATE, above_obj, 18, False, obj), |
| (above_obj, grasp, 12, False, obj), |
| (grasp, close, 8, False, obj), |
| (close, lift, 14, True, obj), |
| (lift, above_goal, 20, True, obj), |
| (above_goal, place, 12, True, obj), |
| (place, open_at_goal, 8, True, goal), |
| (open_at_goal, retreat, 12, False, goal), |
| ] |
|
|
| frames: list[tuple[np.ndarray, bool, np.ndarray]] = [] |
| for start, end, steps, attached, cube_anchor in segments: |
| for state in interpolate(start, end, steps): |
| frames.append((state, attached, cube_anchor.copy())) |
| return frames |
|
|
|
|
| def dataset_features(image_key: str, width: int, height: int) -> dict: |
| return { |
| "action": {"dtype": "float32", "shape": (6,), "names": JOINT_NAMES}, |
| "observation.state": {"dtype": "float32", "shape": (6,), "names": JOINT_NAMES}, |
| f"observation.images.{image_key}": { |
| "dtype": "image", |
| "shape": (height, width, 3), |
| "names": ["height", "width", "channels"], |
| }, |
| } |
|
|
|
|
| def collect_dataset(cfg: AppConfig) -> None: |
| if cfg.dataset.root.exists(): |
| if not cfg.dataset.overwrite: |
| raise FileExistsError(f"Dataset root already exists: {cfg.dataset.root}") |
| shutil.rmtree(cfg.dataset.root) |
|
|
| dataset = LeRobotDataset.create( |
| repo_id=cfg.dataset.repo_id, |
| root=cfg.dataset.root, |
| fps=cfg.dataset.fps, |
| features=dataset_features(cfg.dataset.image_key, cfg.dataset.width, cfg.dataset.height), |
| robot_type=cfg.dataset.robot_type, |
| use_videos=False, |
| ) |
|
|
| model, data, ids = prepare_model(cfg.dataset.base_height) |
| renderer = mujoco.Renderer(model, height=cfg.dataset.height, width=cfg.dataset.width) |
| rng = np.random.default_rng(cfg.dataset.seed) |
| dt = 1.0 / cfg.dataset.fps |
|
|
| if cfg.dataset.show_viewer: |
| viewer_cm = mujoco.viewer.launch_passive(model, data) |
| else: |
| viewer_cm = None |
|
|
| try: |
| if viewer_cm is not None: |
| viewer = viewer_cm.__enter__() |
| configure_viewer_camera(viewer, cfg.dataset.base_height) |
| else: |
| viewer = None |
|
|
| for _ in range(cfg.dataset.episodes): |
| obj_pos, goal_pos = sample_positions(rng) |
| model.site_pos[ids["goal_site"]] = goal_pos |
| frames = make_episode_trajectory(obj_pos, goal_pos) |
| prev_state = frames[0][0].copy() |
| attached_once = False |
|
|
| for state, attached, cube_anchor in frames: |
| start = time.perf_counter() |
| if attached and not attached_once: |
| attached_once = True |
|
|
| set_arm_state(model, data, ids, state) |
| mujoco.mj_forward(model, data) |
| tcp_pos = data.site_xpos[ids["tcp"]].copy() |
| if attached: |
| cube_pos = tcp_pos + ATTACH_OFFSET |
| elif attached_once: |
| cube_pos = goal_pos.copy() |
| else: |
| cube_pos = cube_anchor.copy() |
| set_cube_pose(model, data, ids, cube_pos) |
| mujoco.mj_forward(model, data) |
| image = render_frame(renderer, data, cfg.dataset.base_height) |
|
|
| dataset.add_frame( |
| { |
| "observation.state": prev_state.astype(np.float32), |
| "action": state.astype(np.float32), |
| f"observation.images.{cfg.dataset.image_key}": image, |
| "task": cfg.dataset.task, |
| } |
| ) |
| prev_state = state.copy() |
|
|
| if viewer is not None: |
| viewer.sync() |
| if not viewer.is_running(): |
| viewer = None |
| else: |
| elapsed = time.perf_counter() - start |
| if elapsed < dt: |
| time.sleep(dt - elapsed) |
|
|
| dataset.save_episode() |
| finally: |
| renderer.close() |
| dataset.finalize() |
| if viewer_cm is not None: |
| viewer_cm.__exit__(None, None, None) |
|
|
|
|
| def run_subprocess(cmd: list[str]) -> None: |
| subprocess.run(cmd, check=True, env=lerobot_env()) |
|
|
|
|
| def viz_dataset(cfg: AppConfig) -> None: |
| cmd = [ |
| "lerobot-dataset-viz", |
| "--repo-id", |
| cfg.dataset.repo_id, |
| "--root", |
| str(cfg.dataset.root), |
| "--episode-index", |
| str(cfg.viz.episode_index), |
| ] |
| run_subprocess(cmd) |
|
|
|
|
| def train_policy(cfg: AppConfig) -> None: |
| cmd = [ |
| "lerobot-train", |
| f"--dataset.repo_id={cfg.dataset.repo_id}", |
| f"--dataset.root={cfg.dataset.root}", |
| "--policy.type=act", |
| "--policy.push_to_hub=false", |
| f"--output_dir={cfg.train.output_dir}", |
| f"--job_name={cfg.train.job_name}", |
| f"--policy.device={cfg.train.device}", |
| f"--batch_size={cfg.train.batch_size}", |
| f"--num_workers={cfg.train.num_workers}", |
| f"--steps={cfg.train.steps}", |
| f"--save_freq={cfg.train.save_freq}", |
| f"--log_freq={cfg.train.log_freq}", |
| f"--eval_freq={cfg.train.eval_freq}", |
| f"--seed={cfg.train.seed}", |
| f"--wandb.enable={bool_cli(cfg.train.wandb_enable)}", |
| f"--resume={bool_cli(cfg.train.resume)}", |
| ] |
| print("command=" + " ".join(shlex.quote(part) for part in cmd)) |
| sys.stdout.flush() |
| sys.stderr.flush() |
| os.execvpe(cmd[0], cmd, lerobot_env()) |
|
|
|
|
| def resolve_policy_root(cfg: AppConfig) -> Path: |
| if cfg.eval.policy_root is not None: |
| return cfg.eval.policy_root |
| candidates = sorted(cfg.train.output_dir.glob("checkpoints/*/pretrained_model")) |
| if not candidates: |
| raise FileNotFoundError(f"No checkpoint found under {cfg.train.output_dir / 'checkpoints'}") |
| return candidates[-1] |
|
|
|
|
| def load_policy(policy_root: Path, device: torch.device): |
| policy_cfg = PreTrainedConfig.from_pretrained(policy_root, cli_overrides=[f"--device={device.type}"]) |
| policy_cls = get_policy_class(policy_cfg.type) |
| policy = policy_cls.from_pretrained(policy_root, config=policy_cfg, local_files_only=True) |
| processor_overrides = {"device_processor": {"device": device.type}} |
| preprocessor, postprocessor = make_pre_post_processors( |
| policy_cfg, |
| pretrained_path=str(policy_root), |
| preprocessor_overrides=processor_overrides, |
| postprocessor_overrides=processor_overrides, |
| ) |
| return policy, preprocessor, postprocessor, policy_cfg |
|
|
|
|
| def get_policy_observation_keys(policy_cfg: PreTrainedConfig) -> tuple[str, str]: |
| state_key = None |
| image_key = None |
| for key, feature in policy_cfg.input_features.items(): |
| if feature.type is FeatureType.STATE and state_key is None: |
| state_key = key |
| if feature.type is FeatureType.VISUAL and image_key is None: |
| image_key = key |
| if state_key is None or image_key is None: |
| raise ValueError("Policy must contain one state input and one image input.") |
| return state_key, image_key |
|
|
|
|
| def prepare_policy_input( |
| state: np.ndarray, |
| image: np.ndarray, |
| state_key: str, |
| image_key: str, |
| task: str, |
| robot_type: str, |
| device: torch.device, |
| ) -> dict: |
| return { |
| state_key: torch.from_numpy(state.astype(np.float32)).unsqueeze(0).to(device), |
| image_key: ( |
| torch.from_numpy(image.astype(np.float32) / 255.0) |
| .permute(2, 0, 1) |
| .contiguous() |
| .unsqueeze(0) |
| .to(device) |
| ), |
| "task": task, |
| "robot_type": robot_type, |
| } |
|
|
|
|
| def action_limits(model: mujoco.MjModel, ids: dict[str, int]) -> tuple[np.ndarray, np.ndarray]: |
| lower = [] |
| upper = [] |
| for joint_name in ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll"]: |
| low, high = model.jnt_range[ids[joint_name]] |
| lower.append(np.rad2deg(low)) |
| upper.append(np.rad2deg(high)) |
| lower.append(0.0) |
| upper.append(100.0) |
| return np.asarray(lower, dtype=np.float64), np.asarray(upper, dtype=np.float64) |
|
|
|
|
| def update_cube_attachment( |
| attached: bool, |
| cube_pos: np.ndarray, |
| tcp_pos: np.ndarray, |
| action: np.ndarray, |
| goal_pos: np.ndarray, |
| ) -> tuple[bool, np.ndarray]: |
| grip = float(action[5]) |
| tcp_to_cube = np.linalg.norm(tcp_pos - (cube_pos - ATTACH_OFFSET)) |
| if attached: |
| if grip > GRIP_OPEN_THRESHOLD: |
| released = np.array([tcp_pos[0], tcp_pos[1], TABLE_Z], dtype=np.float64) |
| if np.linalg.norm(released[:2] - goal_pos[:2]) < SUCCESS_XY_THRESHOLD: |
| released[:2] = goal_pos[:2] |
| return False, released |
| return True, tcp_pos + ATTACH_OFFSET |
| if grip < GRIP_CLOSE_THRESHOLD and tcp_to_cube < GRASP_DISTANCE_THRESHOLD: |
| return True, tcp_pos + ATTACH_OFFSET |
| return False, cube_pos |
|
|
|
|
| def check_success(attached: bool, cube_pos: np.ndarray, goal_pos: np.ndarray) -> bool: |
| if attached: |
| return False |
| return np.linalg.norm(cube_pos[:2] - goal_pos[:2]) < SUCCESS_XY_THRESHOLD and abs(cube_pos[2] - TABLE_Z) < 1e-3 |
|
|
|
|
| def eval_policy(cfg: AppConfig) -> None: |
| policy_root = resolve_policy_root(cfg) |
| device = parse_device(cfg.eval.device) |
| policy, preprocessor, postprocessor, policy_cfg = load_policy(policy_root, device) |
| state_key, image_key = get_policy_observation_keys(policy_cfg) |
| image_shape = policy_cfg.input_features[image_key].shape |
| height = int(image_shape[1]) |
| width = int(image_shape[2]) |
|
|
| model, data, ids = prepare_model(cfg.eval.base_height) |
| renderer = mujoco.Renderer(model, height=height, width=width) |
| lower, upper = action_limits(model, ids) |
| rng = np.random.default_rng(cfg.eval.seed) |
| successes = 0 |
|
|
| if cfg.eval.show_viewer: |
| viewer_cm = mujoco.viewer.launch_passive(model, data) |
| else: |
| viewer_cm = None |
|
|
| try: |
| if viewer_cm is not None: |
| viewer = viewer_cm.__enter__() |
| configure_viewer_camera(viewer, cfg.eval.base_height) |
| else: |
| viewer = None |
|
|
| for episode_index in range(cfg.eval.episodes): |
| policy.reset() |
| preprocessor.reset() |
| postprocessor.reset() |
|
|
| object_pos, goal_pos = sample_positions(rng) |
| model.site_pos[ids["goal_site"]] = goal_pos |
| current_state = HOME_STATE.copy() |
| attached = False |
| cube_pos = object_pos.copy() |
| success = False |
|
|
| for _ in range(cfg.eval.max_steps): |
| set_arm_state(model, data, ids, current_state) |
| set_cube_pose(model, data, ids, cube_pos) |
| mujoco.mj_forward(model, data) |
| image = render_frame(renderer, data, cfg.eval.base_height) |
| obs = prepare_policy_input( |
| state=current_state, |
| image=image, |
| state_key=state_key, |
| image_key=image_key, |
| task=cfg.eval.task, |
| robot_type=cfg.eval.robot_type, |
| device=device, |
| ) |
|
|
| with torch.inference_mode(): |
| processed = preprocessor(obs) |
| action = policy.select_action(processed) |
| action = postprocessor(action) |
|
|
| current_state = action.squeeze(0).detach().to("cpu").numpy().astype(np.float64) |
| current_state = np.clip(current_state, lower, upper) |
|
|
| set_arm_state(model, data, ids, current_state) |
| mujoco.mj_forward(model, data) |
| tcp_pos = data.site_xpos[ids["tcp"]].copy() |
| attached, cube_pos = update_cube_attachment(attached, cube_pos, tcp_pos, current_state, goal_pos) |
| set_cube_pose(model, data, ids, cube_pos) |
| mujoco.mj_forward(model, data) |
|
|
| if viewer is not None: |
| viewer.sync() |
| if not viewer.is_running(): |
| viewer = None |
| elif cfg.eval.step_sleep > 0: |
| time.sleep(cfg.eval.step_sleep) |
|
|
| if check_success(attached, cube_pos, goal_pos): |
| success = True |
| break |
|
|
| successes += int(success) |
| print(f"episode={episode_index} success={int(success)} cube={cube_pos.round(4).tolist()} goal={goal_pos.round(4).tolist()}") |
|
|
| print(f"policy_root={policy_root}") |
| print(f"episodes={cfg.eval.episodes}") |
| print(f"successes={successes}") |
| print(f"success_rate={successes / cfg.eval.episodes:.3f}") |
| finally: |
| renderer.close() |
| if viewer_cm is not None: |
| viewer_cm.__exit__(None, None, None) |
|
|
