chore: publish H2 real-world RL system

d8290eb verified 3 days ago

11.9 kB

#!/usr/bin/env python3
"""Offline H2 artifact/mapping/policy/CLoSD preflight; never opens robot topics."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import sys
import time

import numpy as np
import torch


ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from runtime.configuration import load_settings  # noqa: E402
from runtime.embodiments import load_embodiment  # noqa: E402
from runtime.h2_contract import (  # noqa: E402
    EXPECTED_HW_TO_POLICY,
    EXPECTED_POLICY_TO_HW,
    HARDWARE_JOINT_NAMES,
    H2_JOINT_MAP_REVISION,
    H2Contract,
    H2_DOF,
)
from runtime.h2_shm_bridge import BRIDGE_SIZE, SHM_VERSION  # noqa: E402
from runtime.observation import build_observation, expected_obs_dim  # noqa: E402
from runtime.policy import OnnxPolicy  # noqa: E402
from runtime.provenance import verify_artifacts  # noqa: E402
from runtime.deployment_versions import DEPLOYMENTS  # noqa: E402
from runtime.real_reference import RealReferenceState  # noqa: E402
from runtime.reference import ReferenceGenerator, load_frozen_pose  # noqa: E402


def run(config: Path, *, policy_only: bool, deployment_version: str | None = None) -> dict:
    settings = load_settings(config, deployment_version=deployment_version)
    verified = verify_artifacts()
    spec = load_embodiment("h2")
    policy = OnnxPolicy(str(settings.policy_path), device=settings.policy_device)
    contract = H2Contract.from_policy(policy, spec.mjcf_path())
    frozen = load_frozen_pose(spec, list(contract.policy_names))
    ready_hw = contract.ready_hw(frozen.joint_pos)

    probe_hw = np.arange(H2_DOF, dtype=np.float64)
    roundtrip = contract.hw_from_policy(contract.policy_from_hw(probe_hw))
    if not np.array_equal(probe_hw, roundtrip):
        raise RuntimeError("H2 mapping round-trip failed")
    if tuple(contract.hw_to_policy.tolist()) != EXPECTED_HW_TO_POLICY:
        raise RuntimeError("H2 hardware-to-policy mapping tripwire failed")
    if tuple(contract.policy_to_hw.tolist()) != EXPECTED_POLICY_TO_HW:
        raise RuntimeError("H2 policy-to-hardware mapping tripwire failed")
    if settings.hardware_joint_map_revision != H2_JOINT_MAP_REVISION:
        raise RuntimeError(
            "configured H2 joint-map revision does not match the audited executable: "
            f"{settings.hardware_joint_map_revision!r} != {H2_JOINT_MAP_REVISION!r}"
        )
    critical_joint_slots = {
        4: ("left_ankle_roll_joint", 14),
        5: ("left_ankle_pitch_joint", 19),
        10: ("right_ankle_roll_joint", 15),
        11: ("right_ankle_pitch_joint", 20),
        12: ("waist_yaw_joint", 2),
        13: ("waist_roll_joint", 5),
        14: ("waist_pitch_joint", 8),
        19: ("left_wrist_roll_joint", 25),
        20: ("left_wrist_pitch_joint", 27),
        21: ("left_wrist_yaw_joint", 29),
        26: ("right_wrist_roll_joint", 26),
        27: ("right_wrist_pitch_joint", 28),
        28: ("right_wrist_yaw_joint", 30),
    }
    for hardware_id, (joint_name, policy_id) in critical_joint_slots.items():
        if HARDWARE_JOINT_NAMES[hardware_id] != joint_name:
            raise RuntimeError(
                f"H2 joint-map name tripwire failed at hardware id {hardware_id}: "
                f"{HARDWARE_JOINT_NAMES[hardware_id]!r} != {joint_name!r}"
            )
        if int(contract.hw_to_policy[hardware_id]) != policy_id:
            raise RuntimeError(
                f"H2 joint-map policy tripwire failed at hardware id {hardware_id}: "
                f"{int(contract.hw_to_policy[hardware_id])} != {policy_id}"
            )
    if not np.allclose(ready_hw[29:31], 0.0):
        raise RuntimeError("head lock was not applied to official H2 motor ids 29/30")

    reference_joint_pos = frozen.joint_pos.copy()
    reference_joint_vel = np.zeros(H2_DOF, dtype=np.float64)
    reference_pelvis_pos = frozen.pelvis_pos.copy()
    reference_pelvis_quat = frozen.pelvis_quat.copy()
    generation = None
    if not policy_only:
        if settings.generator_device == "cuda" and not torch.cuda.is_available():
            raise RuntimeError("full preflight requested CUDA but torch.cuda.is_available() is false")
        generator = ReferenceGenerator(
            spec,
            list(contract.policy_names),
            contract.default_policy,
            device=settings.generator_device,
            sampling_steps=settings.generator_sampling_steps,
        )
        reference_state = RealReferenceState(
            spec, generator, frozen, landing_xy=settings.landing_xy,
            latency_skip_s=settings.active_latency_skip_s,
        )
        try:
            # Mirror the real runner: warm CUDA kernels, then exercise the
            # asynchronous request/result/latency-skip path with no robot I/O.
            reference_state.warmup(settings.generator_warmup_steps)
            reference_state.update_robot_state(
                np.zeros(2), np.array([1.0, 0.0, 0.0, 0.0]), frozen.joint_pos
            )
            reference_state.update_external_command(
                1, 0.80, np.array([0.40, 0.0, 1.10]), np.array([-3.0, 0.0, 0.2]), base_y=0.0
            )
            deadline = time.monotonic() + 30.0
            while reference_state.command_id != 1 and time.monotonic() < deadline:
                reference_state.step()
                time.sleep(0.005)
            reference_state.raise_if_failed()
            if reference_state.command_id != 1 or reference_state.is_frozen:
                raise RuntimeError("asynchronous H2 CLoSD result did not activate before timeout")
            arrays = [
                reference_state.ref_joint_pos,
                reference_state.ref_joint_vel,
                reference_state.ref_pelvis_pos,
                reference_state.ref_pelvis_quat,
            ]
            if not all(np.isfinite(array).all() for array in arrays):
                raise RuntimeError("asynchronous H2 CLoSD reference contains NaN/Inf")
            phase_values = np.asarray(
                [
                    reference_state.last_motion_hit_time_s,
                    reference_state.last_remaining_command_s,
                ],
                dtype=np.float64,
            )
            if not np.isfinite(phase_values).all():
                raise RuntimeError("H2 CLoSD phase alignment diagnostics are non-finite")
            phase_frames = int(
                round(
                    (
                        reference_state.last_motion_hit_time_s
                        - reference_state.last_remaining_command_s
                    )
                    * generator.fps
                )
            )
            expected_skip = phase_frames if phase_frames >= 1 else 1
            expected_wait = 0 if phase_frames >= 1 else 1 - phase_frames
            if reference_state.last_skip_frames != expected_skip:
                raise RuntimeError(
                    "H2 CLoSD phase skip invariant failed: "
                    f"{reference_state.last_skip_frames} != {expected_skip}"
                )
            if reference_state.last_wait_frames != expected_wait:
                raise RuntimeError(
                    "H2 CLoSD phase wait invariant failed: "
                    f"{reference_state.last_wait_frames} != {expected_wait}"
                )
            if reference_state.last_skip_frames < 1:
                raise RuntimeError("H2 CLoSD must always discard generated frame zero")
            if not np.isclose(
                reference_state.last_motion_hit_time_s,
                spec.condition_time_range[1],
                atol=1.0e-5,
            ):
                raise RuntimeError(
                    "the 0.80s preflight command was not clamped to the H2 bundle's "
                    f"effective hit time {spec.condition_time_range[1]:.3f}s"
                )
            reference_joint_pos = reference_state.ref_joint_pos.copy()
            reference_joint_vel = reference_state.ref_joint_vel.copy()
            reference_pelvis_pos = reference_state.ref_pelvis_pos.copy()
            reference_pelvis_quat = reference_state.ref_pelvis_quat.copy()
            generation = {
                "frames": generator.pred_len,
                "clamped": reference_state.last_clamped,
                "clamp_delta_m": reference_state.last_clamp_delta_m,
                "device": generator.device,
                "generation_s": reference_state.last_generation_s,
                "latency_skip_frames": reference_state.last_skip_frames,
                "skip_frames": reference_state.last_skip_frames,
                "wait_frames": reference_state.last_wait_frames,
                "motion_hit_s": reference_state.last_motion_hit_time_s,
                "remaining_command_s": reference_state.last_remaining_command_s,
                "async_path": True,
            }
        finally:
            reference_state.close()

    obs = build_observation(
        task_command=np.asarray(spec.frozen_command(), dtype=np.float64),
        ref_joint_pos=reference_joint_pos,
        ref_joint_vel=reference_joint_vel,
        root_xy=np.zeros(2),
        ref_pelvis_xy=reference_pelvis_pos[:2],
        ref_pelvis_quat=reference_pelvis_quat,
        pelvis_quat=np.array([1.0, 0.0, 0.0, 0.0]),
        heading_quat=np.array([1.0, 0.0, 0.0, 0.0]),
        base_ang_vel_body=np.zeros(3),
        joint_pos=frozen.joint_pos,
        joint_vel=np.zeros(H2_DOF),
        default_joint_pos=contract.default_policy,
        prev_action=np.zeros(H2_DOF),
        noise=False,
    )
    if obs.shape != (expected_obs_dim(H2_DOF),) or not np.isfinite(obs).all():
        raise RuntimeError(f"invalid H2 observation: {obs.shape}")
    action = np.clip(policy(obs), -10.0, 10.0)
    # Exercise the actual learned-policy contract: the clipped action maps to
    # a raw virtual PD target.  Mechanical position limiting is reserved for
    # ready/stop; the bridge performs the latest-state effort projection.
    target_hw = contract.learned_raw_hw_targets_from_action(action)
    if not np.isfinite(target_hw).all() or not np.allclose(target_hw[29:31], 0.0):
        raise RuntimeError("invalid policy-to-H2 target conversion")

    return {
        "deployment_version": settings.deployment_version,
        "policy_path": str(settings.policy_path),
        "status": "PASS",
        "artifacts_verified": len(verified),
        "policy_obs_dim": policy.obs_dim,
        "policy_action_dim": policy.act_dim,
        "shm_bytes": BRIDGE_SIZE,
        "shm_version": SHM_VERSION,
        "mode_pr": 0,
        "joint_map_revision": H2_JOINT_MAP_REVISION,
        "hw_to_policy": contract.hw_to_policy.tolist(),
        "policy_to_hw": contract.policy_to_hw.tolist(),
        "critical_joint_slots_by_hw_id": {
            str(hardware_id): {
                "name": joint_name,
                "policy_index": policy_id,
            }
            for hardware_id, (joint_name, policy_id) in critical_joint_slots.items()
        },
        "head_target_hw_29_30": target_hw[29:31].tolist(),
        "generation": generation,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=ROOT / "config.yaml")
    parser.add_argument("--deployment-version", "--upstream-version", choices=DEPLOYMENTS)
    parser.add_argument("--policy-only", action="store_true", help="Skip the native CLoSD generation smoke test.")
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()
    report = run(args.config, policy_only=args.policy_only, deployment_version=args.deployment_version)
    if args.json:
        print(json.dumps(report, indent=2, sort_keys=True))
    else:
        print("[PREFLIGHT PASS] " + json.dumps(report, sort_keys=True))


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""Offline H2 artifact/mapping/policy/CLoSD preflight; never opens robot topics."""

	from __future__ import annotations

	import argparse
	import json
	from pathlib import Path
	import sys
	import time

	import numpy as np
	import torch


	ROOT = Path(__file__).resolve().parents[1]
	if str(ROOT) not in sys.path:
	sys.path.insert(0, str(ROOT))

	from runtime.configuration import load_settings # noqa: E402
	from runtime.embodiments import load_embodiment # noqa: E402
	from runtime.h2_contract import ( # noqa: E402
	EXPECTED_HW_TO_POLICY,
	EXPECTED_POLICY_TO_HW,
	HARDWARE_JOINT_NAMES,
	H2_JOINT_MAP_REVISION,
	H2Contract,
	H2_DOF,
	)
	from runtime.h2_shm_bridge import BRIDGE_SIZE, SHM_VERSION # noqa: E402
	from runtime.observation import build_observation, expected_obs_dim # noqa: E402
	from runtime.policy import OnnxPolicy # noqa: E402
	from runtime.provenance import verify_artifacts # noqa: E402
	from runtime.deployment_versions import DEPLOYMENTS # noqa: E402
	from runtime.real_reference import RealReferenceState # noqa: E402
	from runtime.reference import ReferenceGenerator, load_frozen_pose # noqa: E402


	def run(config: Path, *, policy_only: bool, deployment_version: str \| None = None) -> dict:
	settings = load_settings(config, deployment_version=deployment_version)
	verified = verify_artifacts()
	spec = load_embodiment("h2")
	policy = OnnxPolicy(str(settings.policy_path), device=settings.policy_device)
	contract = H2Contract.from_policy(policy, spec.mjcf_path())
	frozen = load_frozen_pose(spec, list(contract.policy_names))
	ready_hw = contract.ready_hw(frozen.joint_pos)

	probe_hw = np.arange(H2_DOF, dtype=np.float64)
	roundtrip = contract.hw_from_policy(contract.policy_from_hw(probe_hw))
	if not np.array_equal(probe_hw, roundtrip):
	raise RuntimeError("H2 mapping round-trip failed")
	if tuple(contract.hw_to_policy.tolist()) != EXPECTED_HW_TO_POLICY:
	raise RuntimeError("H2 hardware-to-policy mapping tripwire failed")
	if tuple(contract.policy_to_hw.tolist()) != EXPECTED_POLICY_TO_HW:
	raise RuntimeError("H2 policy-to-hardware mapping tripwire failed")
	if settings.hardware_joint_map_revision != H2_JOINT_MAP_REVISION:
	raise RuntimeError(
	"configured H2 joint-map revision does not match the audited executable: "
	f"{settings.hardware_joint_map_revision!r} != {H2_JOINT_MAP_REVISION!r}"
	)
	critical_joint_slots = {
	4: ("left_ankle_roll_joint", 14),
	5: ("left_ankle_pitch_joint", 19),
	10: ("right_ankle_roll_joint", 15),
	11: ("right_ankle_pitch_joint", 20),
	12: ("waist_yaw_joint", 2),
	13: ("waist_roll_joint", 5),
	14: ("waist_pitch_joint", 8),
	19: ("left_wrist_roll_joint", 25),
	20: ("left_wrist_pitch_joint", 27),
	21: ("left_wrist_yaw_joint", 29),
	26: ("right_wrist_roll_joint", 26),
	27: ("right_wrist_pitch_joint", 28),
	28: ("right_wrist_yaw_joint", 30),
	}
	for hardware_id, (joint_name, policy_id) in critical_joint_slots.items():
	if HARDWARE_JOINT_NAMES[hardware_id] != joint_name:
	raise RuntimeError(
	f"H2 joint-map name tripwire failed at hardware id {hardware_id}: "
	f"{HARDWARE_JOINT_NAMES[hardware_id]!r} != {joint_name!r}"
	)
	if int(contract.hw_to_policy[hardware_id]) != policy_id:
	raise RuntimeError(
	f"H2 joint-map policy tripwire failed at hardware id {hardware_id}: "
	f"{int(contract.hw_to_policy[hardware_id])} != {policy_id}"
	)
	if not np.allclose(ready_hw[29:31], 0.0):
	raise RuntimeError("head lock was not applied to official H2 motor ids 29/30")

	reference_joint_pos = frozen.joint_pos.copy()
	reference_joint_vel = np.zeros(H2_DOF, dtype=np.float64)
	reference_pelvis_pos = frozen.pelvis_pos.copy()
	reference_pelvis_quat = frozen.pelvis_quat.copy()
	generation = None
	if not policy_only:
	if settings.generator_device == "cuda" and not torch.cuda.is_available():
	raise RuntimeError("full preflight requested CUDA but torch.cuda.is_available() is false")
	generator = ReferenceGenerator(
	spec,
	list(contract.policy_names),
	contract.default_policy,
	device=settings.generator_device,
	sampling_steps=settings.generator_sampling_steps,
	)
	reference_state = RealReferenceState(
	spec, generator, frozen, landing_xy=settings.landing_xy,
	latency_skip_s=settings.active_latency_skip_s,
	)
	try:
	# Mirror the real runner: warm CUDA kernels, then exercise the
	# asynchronous request/result/latency-skip path with no robot I/O.
	reference_state.warmup(settings.generator_warmup_steps)
	reference_state.update_robot_state(
	np.zeros(2), np.array([1.0, 0.0, 0.0, 0.0]), frozen.joint_pos
	)
	reference_state.update_external_command(
	1, 0.80, np.array([0.40, 0.0, 1.10]), np.array([-3.0, 0.0, 0.2]), base_y=0.0
	)
	deadline = time.monotonic() + 30.0
	while reference_state.command_id != 1 and time.monotonic() < deadline:
	reference_state.step()
	time.sleep(0.005)
	reference_state.raise_if_failed()
	if reference_state.command_id != 1 or reference_state.is_frozen:
	raise RuntimeError("asynchronous H2 CLoSD result did not activate before timeout")
	arrays = [
	reference_state.ref_joint_pos,
	reference_state.ref_joint_vel,
	reference_state.ref_pelvis_pos,
	reference_state.ref_pelvis_quat,
	]
	if not all(np.isfinite(array).all() for array in arrays):
	raise RuntimeError("asynchronous H2 CLoSD reference contains NaN/Inf")
	phase_values = np.asarray(
	[
	reference_state.last_motion_hit_time_s,
	reference_state.last_remaining_command_s,
	],
	dtype=np.float64,
	)
	if not np.isfinite(phase_values).all():
	raise RuntimeError("H2 CLoSD phase alignment diagnostics are non-finite")
	phase_frames = int(
	round(
	(
	reference_state.last_motion_hit_time_s
	- reference_state.last_remaining_command_s
	)
	* generator.fps
	)
	)
	expected_skip = phase_frames if phase_frames >= 1 else 1
	expected_wait = 0 if phase_frames >= 1 else 1 - phase_frames
	if reference_state.last_skip_frames != expected_skip:
	raise RuntimeError(
	"H2 CLoSD phase skip invariant failed: "
	f"{reference_state.last_skip_frames} != {expected_skip}"
	)
	if reference_state.last_wait_frames != expected_wait:
	raise RuntimeError(
	"H2 CLoSD phase wait invariant failed: "
	f"{reference_state.last_wait_frames} != {expected_wait}"
	)
	if reference_state.last_skip_frames < 1:
	raise RuntimeError("H2 CLoSD must always discard generated frame zero")
	if not np.isclose(
	reference_state.last_motion_hit_time_s,
	spec.condition_time_range[1],
	atol=1.0e-5,
	):
	raise RuntimeError(
	"the 0.80s preflight command was not clamped to the H2 bundle's "
	f"effective hit time {spec.condition_time_range[1]:.3f}s"
	)
	reference_joint_pos = reference_state.ref_joint_pos.copy()
	reference_joint_vel = reference_state.ref_joint_vel.copy()
	reference_pelvis_pos = reference_state.ref_pelvis_pos.copy()
	reference_pelvis_quat = reference_state.ref_pelvis_quat.copy()
	generation = {
	"frames": generator.pred_len,
	"clamped": reference_state.last_clamped,
	"clamp_delta_m": reference_state.last_clamp_delta_m,
	"device": generator.device,
	"generation_s": reference_state.last_generation_s,
	"latency_skip_frames": reference_state.last_skip_frames,
	"skip_frames": reference_state.last_skip_frames,
	"wait_frames": reference_state.last_wait_frames,
	"motion_hit_s": reference_state.last_motion_hit_time_s,
	"remaining_command_s": reference_state.last_remaining_command_s,
	"async_path": True,
	}
	finally:
	reference_state.close()

	obs = build_observation(
	task_command=np.asarray(spec.frozen_command(), dtype=np.float64),
	ref_joint_pos=reference_joint_pos,
	ref_joint_vel=reference_joint_vel,
	root_xy=np.zeros(2),
	ref_pelvis_xy=reference_pelvis_pos[:2],
	ref_pelvis_quat=reference_pelvis_quat,
	pelvis_quat=np.array([1.0, 0.0, 0.0, 0.0]),
	heading_quat=np.array([1.0, 0.0, 0.0, 0.0]),
	base_ang_vel_body=np.zeros(3),
	joint_pos=frozen.joint_pos,
	joint_vel=np.zeros(H2_DOF),
	default_joint_pos=contract.default_policy,
	prev_action=np.zeros(H2_DOF),
	noise=False,
	)
	if obs.shape != (expected_obs_dim(H2_DOF),) or not np.isfinite(obs).all():
	raise RuntimeError(f"invalid H2 observation: {obs.shape}")
	action = np.clip(policy(obs), -10.0, 10.0)
	# Exercise the actual learned-policy contract: the clipped action maps to
	# a raw virtual PD target. Mechanical position limiting is reserved for
	# ready/stop; the bridge performs the latest-state effort projection.
	target_hw = contract.learned_raw_hw_targets_from_action(action)
	if not np.isfinite(target_hw).all() or not np.allclose(target_hw[29:31], 0.0):
	raise RuntimeError("invalid policy-to-H2 target conversion")

	return {
	"deployment_version": settings.deployment_version,
	"policy_path": str(settings.policy_path),
	"status": "PASS",
	"artifacts_verified": len(verified),
	"policy_obs_dim": policy.obs_dim,
	"policy_action_dim": policy.act_dim,
	"shm_bytes": BRIDGE_SIZE,
	"shm_version": SHM_VERSION,
	"mode_pr": 0,
	"joint_map_revision": H2_JOINT_MAP_REVISION,
	"hw_to_policy": contract.hw_to_policy.tolist(),
	"policy_to_hw": contract.policy_to_hw.tolist(),
	"critical_joint_slots_by_hw_id": {
	str(hardware_id): {
	"name": joint_name,
	"policy_index": policy_id,
	}
	for hardware_id, (joint_name, policy_id) in critical_joint_slots.items()
	},
	"head_target_hw_29_30": target_hw[29:31].tolist(),
	"generation": generation,
	}


	def main() -> None:
	parser = argparse.ArgumentParser(description=__doc__)
	parser.add_argument("--config", type=Path, default=ROOT / "config.yaml")
	parser.add_argument("--deployment-version", "--upstream-version", choices=DEPLOYMENTS)
	parser.add_argument("--policy-only", action="store_true", help="Skip the native CLoSD generation smoke test.")
	parser.add_argument("--json", action="store_true")
	args = parser.parse_args()
	report = run(args.config, policy_only=args.policy_only, deployment_version=args.deployment_version)
	if args.json:
	print(json.dumps(report, indent=2, sort_keys=True))
	else:
	print("[PREFLIGHT PASS] " + json.dumps(report, sort_keys=True))


	if __name__ == "__main__":
	main()