chore: publish H2 real-world RL system

d8290eb verified 3 days ago

7.54 kB

"""Per-robot deployment specifications (G1 / H2 / A3).
Every number here mirrors the training configuration of the corresponding
``PingPong-One-Step-<ROBOT>-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 <Robot>Env + gain/armature helpers
    env_class: str
    armature_helper: str | None    # optional build_<robot>_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]

	"""Per-robot deployment specifications (G1 / H2 / A3).

	Every number here mirrors the training configuration of the corresponding
	``PingPong-One-Step-<ROBOT>-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 <Robot>Env + gain/armature helpers
	env_class: str
	armature_helper: str \| None # optional build_<robot>_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]