| |
| """Export the trusted H2 residual actor; preserve the existing collector lifecycle. |
| |
| This is NOT a standalone 182-to-31 tracking policy. Inputs are the original |
| 182 observations followed by the clipped 31 Stage-C actions. Normalization is |
| inside the graph; its seven outputs are right-arm residuals in [-1, 1]. |
| The existing collector still decides WHEN to apply them and combines them in |
| joint radians using ResidualActionAdapter. No learner or exploration is used. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import pickle |
| from pathlib import Path |
| import sys |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
|
|
| class NormalizedResidual(nn.Module): |
| def __init__(self, actor, normalizer, clip): |
| super().__init__() |
| self.actor = actor |
| self.register_buffer('mean', torch.as_tensor(normalizer['mean'], dtype=torch.float32)) |
| self.register_buffer('std', torch.as_tensor(normalizer['std'], dtype=torch.float32)) |
| self.clip = float(clip) |
|
|
| def forward(self, inputs): |
| return self.actor(torch.clamp((inputs - self.mean) / self.std, -self.clip, self.clip)) |
|
|
|
|
| def trusted_load(path): |
| |
| |
| try: |
| return torch.load(path, map_location='cpu', weights_only=True), True |
| except pickle.UnpicklingError: |
| return torch.load(path, map_location='cpu', weights_only=False), False |
|
|
|
|
| def serializable(value): |
| if isinstance(value, np.ndarray): |
| return value.tolist() |
| if isinstance(value, np.generic): |
| return value.item() |
| if isinstance(value, dict): |
| return {key: serializable(item) for key, item in value.items()} |
| if isinstance(value, (tuple, list)): |
| return [serializable(item) for item in value] |
| return value |
|
|
|
|
| def export_checkpoint(root, checkpoint, base_policy, output): |
| root, checkpoint, base_policy, output = map(lambda p: Path(p).resolve(), |
| (root, checkpoint, base_policy, output)) |
| if output.exists() or output.with_suffix('.json').exists(): |
| raise ValueError('choose a new output path; existing models/reports are never replaced') |
| import onnx |
| import onnxruntime as ort |
| sys.path.insert(0, str(root)) |
| from runtime.rl.models import Actor |
| from runtime.rl.data import FORMAT, REWARD_PROTOCOL, ActionContract |
| from runtime.rl.config import RLConfig |
| from runtime.rl.action import ResidualActionAdapter |
|
|
| torch.set_num_threads(1) |
| saved, weights_only = trusted_load(checkpoint) |
| if saved.get('format') != FORMAT or saved.get('reward_protocol') != REWARD_PROTOCOL: |
| raise ValueError('requires the existing H2 real residual v3 checkpoint') |
| config = RLConfig(**saved['config']) |
| recorded = saved['contract'] |
| contract = ActionContract(recorded['policy_names'], recorded['action_scale'], recorded['residual_scales_rad']) |
| contract.validate_saved(recorded) |
| options = ort.SessionOptions() |
| options.intra_op_num_threads = options.inter_op_num_threads = 1 |
| base_session = ort.InferenceSession(str(base_policy), sess_options=options, providers=['CPUExecutionProvider']) |
| metadata = base_session.get_modelmeta().custom_metadata_map |
| base_contract = ActionContract(metadata['joint_names'].split(','), |
| tuple(map(float, metadata['action_scale'].split(','))), config.action_scales_rad) |
| base_contract.validate_saved(recorded) |
| actor = Actor().cpu().eval() |
| actor.load_state_dict(saved['actor']) |
| normalizer = saved['normalizer'] |
| mean = np.asarray(normalizer['mean'], dtype=np.float32) |
| std = np.asarray(normalizer['std'], dtype=np.float32) |
| if (mean.shape != (213,) or std.shape != (213,) or not np.isfinite(mean).all() |
| or not np.isfinite(std).all() or np.any(std <= 0) or not normalizer['frozen']): |
| raise ValueError('requires the saved frozen finite 213-D normalization') |
|
|
| |
| |
| replay = saved['replay']['arrays'] |
| base_checks = [] |
| for obs_key, action_key in (('base_obs', 'base_action'), ('next_base_obs', 'next_base_action')): |
| values = np.asarray(replay[obs_key], dtype=np.float32) |
| inferred = np.concatenate([base_session.run(['actions'], {'obs': row[None]})[0] for row in values]) |
| inferred = np.clip(inferred, -10, 10) |
| recorded_actions = np.asarray(replay[action_key], dtype=np.float32) |
| maximum = float(np.max(np.abs(inferred - recorded_actions))) |
| if not np.array_equal(inferred, recorded_actions): |
| raise ValueError(f'{base_policy.name} differs from recorded {action_key}: max_abs={maximum}') |
| base_checks.append(dict(observations=obs_key, rows=len(values), max_abs_error=maximum, |
| bit_identical=True)) |
|
|
| model = NormalizedResidual(actor, normalizer, config.normalizer_clip).eval() |
| output.parent.mkdir(parents=True, exist_ok=True) |
| torch.onnx.export(model, torch.zeros(1, 213, dtype=torch.float32), str(output), |
| input_names=['inputs'], output_names=['residual'], |
| dynamic_axes={'inputs': {0: 'batch'}, 'residual': {0: 'batch'}}, |
| opset_version=17, dynamo=False) |
| graph = onnx.load(output) |
| onnx.helper.set_model_props(graph, dict(format=FORMAT, component='residual', |
| normalization='embedded', input_layout='obs182,clipped_base_action31', |
| output_layout=','.join(recorded['right_arm_joints']), action_combination='external_existing_ResidualActionAdapter', |
| episode_gate='external_existing_ResidualCollector', source_checkpoint=str(checkpoint), |
| base_policy=str(base_policy))) |
| onnx.checker.check_model(graph) |
| onnx.save(graph, output) |
| session = ort.InferenceSession(str(output), sess_options=options, providers=['CPUExecutionProvider']) |
| rng = np.random.default_rng(20260916) |
| cases = [('replay', np.concatenate((replay['base_obs'], np.clip(replay['base_action'], -10, 10)), axis=1)), |
| ('next_replay', np.concatenate((replay['next_base_obs'], np.clip(replay['next_base_action'], -10, 10)), axis=1)), |
| ('zero', np.zeros((1, 213), np.float32)), |
| ('normalizer_boundaries', np.concatenate((mean[None] + std[None] * config.normalizer_clip, |
| mean[None] - std[None] * config.normalizer_clip))), |
| ('random_extremes', rng.normal(0, 20, (257, 213)).astype(np.float32))] |
| adapter = ResidualActionAdapter(contract) |
| checks = [] |
| for name, inputs in cases: |
| inputs = np.asarray(inputs, dtype=np.float32) |
| |
| normalized = np.clip((inputs - mean) / std, -config.normalizer_clip, config.normalizer_clip) |
| with torch.inference_mode(): |
| expected = actor(torch.from_numpy(normalized)).numpy() |
| actual = session.run(['residual'], {'inputs': inputs})[0] |
| np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=2e-6) |
| if actual.shape != (len(inputs), 7) or not np.isfinite(actual).all() or np.max(np.abs(actual)) > 1: |
| raise ValueError('export must return finite tanh residuals with shape [batch,7]') |
| expected_combined = np.stack([adapter.combine(x[182:], r)[0] for x, r in zip(inputs, expected)]) |
| actual_combined = np.stack([adapter.combine(x[182:], r)[0] for x, r in zip(inputs, actual)]) |
| np.testing.assert_allclose(actual_combined, expected_combined, rtol=1e-5, atol=4e-6) |
| checks.append(dict(case=name, rows=len(inputs), |
| residual_max_abs_error=float(np.max(np.abs(actual - expected))), |
| combined_action_max_abs_error=float(np.max(np.abs(actual_combined - expected_combined))))) |
| report = dict(format=FORMAT, reward_protocol=REWARD_PROTOCOL, |
| config=saved['config'], contract=recorded, counters=saved['counters'], normalizer=normalizer, |
| export=dict(input='raw_concat(obs182,clip(base31,-10,10))', input_name='inputs', output_name='residual', |
| input_dim=213, output_dim=7, normalization='embedded', actor='actor (not actor_target)', |
| source_checkpoint=str(checkpoint), base_policy=str(base_policy), output_path=str(output), |
| checkpoint_sha256=hashlib.sha256(checkpoint.read_bytes()).hexdigest(), |
| base_policy_sha256=hashlib.sha256(base_policy.read_bytes()).hexdigest(), |
| onnx_sha256=hashlib.sha256(output.read_bytes()).hexdigest(), |
| weights_only_load=weights_only, trusted_checkpoint_fallback=not weights_only, |
| base_identity_evidence='All recorded base-action computations match; historical base hash was not stored.', |
| base_replay_checks=base_checks, numerical_checks=checks, |
| standalone_tracking_policy=False, external_episode_gate='unchanged ResidualCollector', |
| external_combination='unchanged ResidualActionAdapter', training=False, exploration=False, |
| hardware_io=False, hardware_qualified=False)) |
| output.with_suffix('.json').write_text(json.dumps(serializable(report), indent=2) + '\n') |
| return serializable(report) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--root', type=Path, required=True) |
| parser.add_argument('--checkpoint', type=Path, required=True) |
| parser.add_argument('--base-policy', type=Path, required=True) |
| parser.add_argument('--output', type=Path, required=True) |
| args = parser.parse_args() |
| print(json.dumps(export_checkpoint(args.root, args.checkpoint, args.base_policy, args.output)['export'], indent=2)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
