"""Per-robot deployment specifications (G1 / H2 / A3). Every number here mirrors the training configuration of the corresponding ``PingPong-One-Step--FB-OnlineCLoSD-Ball-DualCritic`` task: * base task: ``tasks/tracking/tracking_pingpong_online_closd_ball_env_cfg_FB.py`` * robot overrides: ``tasks/tracking/config/{h2,a3}/flat_env_cfg.py`` via ``configure_online_closd_robot``; G1 uses the base values. """ from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path from .paths import assets_dir, package_dir, resolve # The 14 reference bodies tracked by the policy (tracking_pingpong_fk_ref_env_cfg_FB.REFERENCE_BODY_NAMES). REFERENCE_BODY_NAMES: tuple[str, ...] = ( "pelvis", "left_hip_roll_link", "left_knee_link", "left_ankle_roll_link", "right_hip_roll_link", "right_knee_link", "right_ankle_roll_link", "torso_link", "left_shoulder_roll_link", "left_elbow_link", "left_wrist_yaw_link", "right_shoulder_roll_link", "right_elbow_link", "right_wrist_yaw_link", ) @dataclass(frozen=True) class EmbodimentSpec: name: str num_actions: int sim2sim_module: str # module in sim2sim_tt providing Env + gain/armature helpers env_class: str armature_helper: str | None # optional build__armature_for_joint_names in that module mjcf: str # relative to assets/robot_description closd_bundle: str # bundle-relative ready_pose: str # bundle-relative policy: str # bundle-relative default policy paddle_body_name: str paddle_offset_xyz: tuple[float, float, float] paddle_offset_rpy: tuple[float, float, float] paddle_half_width: float paddle_half_height: float nominal_root_z: float hit_x_range: tuple[float, float] hit_y_range: tuple[float, float] hit_z_range: tuple[float, float] condition_time_range: tuple[float, float] condition_lateral_range: tuple[float, float] condition_forward_range: tuple[float, float] condition_height_range: tuple[float, float] foot_body_names: tuple[str, str] # The training articulation has no joint damping / Coulomb friction beyond the # PD actuator model. The MJCFs declare some (A3: 0.5-1.5 N·m·s/rad and # 0.4-1.7 N·m on every joint; H2 0.05 / 0.2; G1 friction 0.1), so they are # neutralised by default; ``--keep-mjcf-joint-dynamics`` keeps them. zero_mjcf_joint_damping: bool = True zero_mjcf_joint_frictionloss: bool = True # shared by all robots (base config) hit_time_range: tuple[float, float] = (0.45, 0.55) post_hit_resample_delay_range: tuple[float, float] = (0.15, 0.70) long_pause_prob: float = 0.2 long_pause_range: tuple[float, float] = (1.5, 3.0) target_landing_x_range: tuple[float, float] = (2.30, 3.10) target_landing_y_range: tuple[float, float] = (-0.60, 0.60) max_delta_hit_y: float = 0.7 stroke_split_relative_y: float = -0.125 frozen_blend_time_s: float = 0.24 frozen_exit_blend_time_s: float = 0.24 closd_fps: float = 50.0 control_dt: float = 0.02 sim_dt: float = 0.005 reference_body_names: tuple[str, ...] = REFERENCE_BODY_NAMES frozen_task_command: tuple[float, ...] = field(default=(-0.59, 0.40, 0.0, 1.0, -3.0, 0.0, 0.2, 2.87, 0.0)) @property def hit_plane_x(self) -> float: return 0.5 * (self.hit_x_range[0] + self.hit_x_range[1]) def mjcf_path(self) -> Path: return resolve(self.mjcf) if Path(self.mjcf).is_absolute() or (package_dir() / self.mjcf).exists() else assets_dir() / self.mjcf def bundle_path(self) -> Path: return resolve(self.closd_bundle) def ready_pose_path(self) -> Path: return resolve(self.ready_pose) def policy_path(self) -> Path: return resolve(self.policy) def frozen_command(self) -> tuple[float, ...]: cmd = list(self.frozen_task_command) cmd[1] = self.hit_plane_x # configure_online_closd_robot rewrites element 1 to hit_plane_x return tuple(cmd) EMBODIMENTS: dict[str, EmbodimentSpec] = { "g1": EmbodimentSpec( name="g1", num_actions=29, sim2sim_module="play_tt_mujoco_g1_fb", env_class="G1Env", armature_helper=None, mjcf="g1/mjcf/g1_full_mesh_with_hand_tt.xml", closd_bundle="closd/G1/model_80000.closd-runtime.pt", ready_pose="ready_pose/G1/frozen_closd_ready_pose_toeout_lowload.npz", policy="policies/g1/policy.onnx", paddle_body_name="right_wrist_yaw_link", paddle_offset_xyz=(0.18522, 0.033288, 0.089488), paddle_offset_rpy=(3.1415926, -0.524, -3.1415926), paddle_half_width=0.0763, paddle_half_height=0.0763, nominal_root_z=0.78, hit_x_range=(0.35, 0.35), hit_y_range=(-0.76, 0.76), hit_z_range=(0.70, 1.30), condition_time_range=(0.40, 0.60), condition_lateral_range=(-1.10, 0.45), condition_forward_range=(-0.05, 0.90), condition_height_range=(0.70, 1.12), foot_body_names=("left_ankle_roll_link", "right_ankle_roll_link"), ), "h2": EmbodimentSpec( name="h2", num_actions=31, sim2sim_module="play_tt_mujoco_h2_fb", env_class="H2Env", armature_helper="build_h2_armature_for_joint_names", mjcf="h2/mjcf/h2_pingpong.xml", closd_bundle="closd/H2/model_80000.closd-runtime.pt", ready_pose="ready_pose/H2/frozen_closd_ready_pose_toeout_lowload.npz", policy="policies/h2/h2_stageC_model_10500.onnx", paddle_body_name="right_wrist_yaw_link", paddle_offset_xyz=(0.19572, 0.036449, 0.08931579), paddle_offset_rpy=(0.0, 0.0, 0.0), paddle_half_width=0.0815, paddle_half_height=0.0855, nominal_root_z=1.0, hit_x_range=(0.35, 0.45), hit_y_range=(-0.76, 0.76), hit_z_range=(0.80, 1.50), condition_time_range=(0.40, 0.60), condition_lateral_range=(-1.35, 0.47), condition_forward_range=(-0.08, 1.16), condition_height_range=(0.99, 1.45), foot_body_names=("left_ankle_pitch_link", "right_ankle_pitch_link"), ), "a3": EmbodimentSpec( name="a3", num_actions=29, sim2sim_module="play_tt_mujoco_a3_fb", env_class="A3Env", armature_helper="build_a3_armature_for_joint_names", mjcf="a3/mjcf/a3_t2d5_pingpong.xml", closd_bundle="closd/A3/model_80000.closd-runtime.pt", ready_pose="ready_pose/A3/frozen_closd_ready_pose_toeout_lowload.npz", policy="policies/a3/policy.onnx", paddle_body_name="right_wrist_yaw_link", paddle_offset_xyz=(0.21021, 0.032078, 0.032036), paddle_offset_rpy=(0.0, 0.0, 0.0), paddle_half_width=0.0763, paddle_half_height=0.0763, nominal_root_z=1.0, hit_x_range=(0.35, 0.45), hit_y_range=(-0.76, 0.76), hit_z_range=(0.70, 1.40), condition_time_range=(0.40, 0.60), condition_lateral_range=(-1.35, 0.40), condition_forward_range=(-0.10, 1.10), condition_height_range=(0.78, 1.26), foot_body_names=("left_ankle_roll_link", "right_ankle_roll_link"), ), } def load_embodiment(name: str) -> EmbodimentSpec: key = name.lower() if key not in EMBODIMENTS: raise KeyError(f"unknown embodiment '{name}', expected one of {sorted(EMBODIMENTS)}") return EMBODIMENTS[key]