| """Versioned H2 contracts; G1 checkpoints cannot be loaded into this runtime.""" |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| import numpy as np |
|
|
| from ..h2_contract import EXPECTED_POLICY_JOINT_NAMES, H2_JOINT_MAP_REVISION |
|
|
| OBS_DIM = 182 |
| BASE_ACTION_DIM = 31 |
| RESIDUAL_ACTION_DIM = 7 |
| ACTOR_INPUT_DIM = OBS_DIM + BASE_ACTION_DIM |
| FORMAT = "h2-real-residual-td3-v3" |
| REWARD_PROTOCOL = "h2-gap-return-n5-v1" |
| RIGHT_ARM_JOINTS = tuple(f"right_{name}_joint" for name in ( |
| "shoulder_pitch", "shoulder_roll", "shoulder_yaw", "elbow", "wrist_roll", "wrist_pitch", "wrist_yaw")) |
|
|
|
|
| def vector(value, size): |
| array = np.asarray(value, dtype=np.float32) |
| if array.shape != (size,) or not np.isfinite(array).all(): |
| raise ValueError(f"expected finite vector ({size},), got {array.shape}") |
| return array.copy() |
|
|
|
|
| @dataclass(frozen=True) |
| class ActionContract: |
| policy_names: tuple[str, ...] |
| action_scale: tuple[float, ...] |
| residual_scales_rad: tuple[float, ...] |
|
|
| def __post_init__(self): |
| for name in ("policy_names", "action_scale", "residual_scales_rad"): |
| object.__setattr__(self, name, tuple(getattr(self, name))) |
| for name in ("action_scale", "residual_scales_rad"): |
| object.__setattr__(self, name, tuple(float(v) for v in getattr(self, name))) |
| if self.policy_names != EXPECTED_POLICY_JOINT_NAMES: |
| raise ValueError("residual policy requires the audited H2 policy joint order") |
| scales = vector(self.action_scale, BASE_ACTION_DIM) |
| residual = vector(self.residual_scales_rad, RESIDUAL_ACTION_DIM) |
| if np.any(scales <= 0) or np.any(residual <= 0): |
| raise ValueError("action scales must be positive") |
|
|
| @property |
| def indices(self): |
| return np.array([self.policy_names.index(n) for n in RIGHT_ARM_JOINTS]) |
|
|
| def to_dict(self): |
| return dict(format=FORMAT, robot="h2", obs_dim=OBS_DIM, action_dim=BASE_ACTION_DIM, |
| residual_dim=RESIDUAL_ACTION_DIM, joint_map_revision=H2_JOINT_MAP_REVISION, |
| policy_names=list(self.policy_names), action_scale=list(self.action_scale), |
| residual_scales_rad=list(self.residual_scales_rad), |
| right_arm_joints=list(RIGHT_ARM_JOINTS), action_clip=10.0) |
|
|
| def validate_saved(self, saved): |
| if saved != self.to_dict(): |
| raise ValueError("checkpoint/action contract mismatch (requires exact H2 dimensions, joint map and scales)") |
|
|
|
|
| @dataclass(frozen=True) |
| class Transition: |
| start_s: float |
| end_s: float |
| base_obs: np.ndarray |
| base_action: np.ndarray |
| res_action: np.ndarray |
| combined_action: np.ndarray |
| next_base_obs: np.ndarray |
| next_base_action: np.ndarray |
| reward: float = 0.0 |
| terminated: bool = False |
| position_distance_m: float | None = None |
| position_reward_missing: bool = False |
|
|
| def __post_init__(self): |
| for name, size in (("base_obs", OBS_DIM), ("next_base_obs", OBS_DIM), |
| ("base_action", BASE_ACTION_DIM), ("next_base_action", BASE_ACTION_DIM), |
| ("combined_action", BASE_ACTION_DIM), ("res_action", RESIDUAL_ACTION_DIM)): |
| object.__setattr__(self, name, vector(getattr(self, name), size)) |
| if not np.isfinite([self.start_s, self.end_s, self.reward]).all() or self.end_s <= self.start_s: |
| raise ValueError("invalid transition interval/reward") |
|
|
|
|
| @dataclass(frozen=True) |
| class EvaluationTransition(Transition): |
| """Measured action interval for scoring, never a residual-training sample.""" |
| observation_dim: int = OBS_DIM |
|
|
| def __post_init__(self): |
| if not isinstance(self.observation_dim, int) or self.observation_dim <= 0: |
| raise ValueError("evaluation observation_dim must be a positive integer") |
| for name, size in (("base_obs", self.observation_dim), ("next_base_obs", self.observation_dim), |
| ("base_action", BASE_ACTION_DIM), ("next_base_action", BASE_ACTION_DIM), |
| ("combined_action", BASE_ACTION_DIM), ("res_action", RESIDUAL_ACTION_DIM)): |
| object.__setattr__(self, name, vector(getattr(self, name), size)) |
| if not np.isfinite([self.start_s, self.end_s, self.reward]).all() or self.end_s <= self.start_s: |
| raise ValueError("invalid transition interval/reward") |
|
|
|
|
| @dataclass(frozen=True) |
| class EpisodeBatch: |
| sequence: int |
| command_id: int |
| phase: str |
| policy_version: int |
| transitions: tuple[Transition, ...] |
| summary: dict |
| collector_rng_state: dict |
|
|
| def validate(self): |
| if self.sequence <= 0 or self.command_id <= 0 or self.phase not in ("warmup", "online"): |
| raise ValueError("invalid episode identity/phase") |
| if not self.transitions or not self.transitions[-1].terminated: |
| raise ValueError("only complete, adjudicated episodes may enter the learner") |
| for i, item in enumerate(self.transitions): |
| if i < len(self.transitions) - 1 and item.terminated: |
| raise ValueError("termination must occur exactly once at the end") |
| if i and abs(self.transitions[i-1].end_s - item.start_s) > 1e-6: |
| raise ValueError("episode transitions must be contiguous") |
|
|
