File size: 5,310 Bytes
d8290eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """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")
|