#!/usr/bin/env python3 """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): # The caller explicitly supplies a trusted local training artifact. Try the # restrictive loader first; this v3 checkpoint also stores NumPy RNG/replay. 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') # The checkpoint did not record a base-model hash. Compare every stored # observation/action pair instead, rather than inventing historical identity. 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) # This is the exact evaluation-only original collector calculation. 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()