"""Validated Go2 velocity actor warm-start; never an optimizer/iteration resume. Callers supply administrator-owned allowed roots, not roots from an HTTP request. Legacy YAML tags are decoded as inert data: no import, eval or object construction. """ import hashlib import io import json from dataclasses import asdict, dataclass from enum import Enum from importlib.metadata import version from pathlib import Path import torch import yaml BASE_TERMS = [ "base_ang_vel", "projected_gravity", "command", "phase", "joint_pos", "joint_vel", "actions", ] JOINTS = [ f"{leg}_{joint}_joint" for leg in ("FL", "FR", "RL", "RR") for joint in ("hip", "thigh", "calf") ] NORMALIZATION_POLICY = "preserve-source-count/unit-new-features" class PretrainedError(ValueError): """The supplied source does not prove the supported transfer contract.""" class _DataLoader(yaml.SafeLoader): def construct_mapping(self, node, deep=False): keys = [self.construct_object(key, deep=deep) for key, _ in node.value] if any(type(key) not in (str, int) for key in keys) or len(set(keys)) != len(keys): raise PretrainedError("Configuration mappings require unique string/integer keys") return super().construct_mapping(node, deep=deep) def _symbol(loader, suffix, node): if loader.construct_scalar(node) != "": raise PretrainedError("Python name tags must have an empty value") return {"symbol": suffix} def _enum(loader, node): values = loader.construct_sequence(node) if len(values) != 1 or not isinstance(values[0], (str, int)): raise PretrainedError("Invalid legacy enum data") return values[0] _DataLoader.add_constructor("tag:yaml.org,2002:python/tuple", _DataLoader.construct_yaml_seq) _DataLoader.add_multi_constructor("tag:yaml.org,2002:python/name:", _symbol) for _name in ("mjlab.actuator.actuator.TransmissionType", "mjlab.viewer.viewer_config.OriginType"): _DataLoader.add_constructor("tag:yaml.org,2002:python/object/apply:" + _name, _enum) _DataLoader.add_constructor( "tag:yaml.org,2002:python/object/apply:builtins.slice", lambda loader, node: {"slice": loader.construct_sequence(node)}, ) def _plain(value): if isinstance(value, Enum): return value.value if callable(value): return {"symbol": value.__module__ + "." + value.__qualname__} if isinstance(value, dict): return {k: _plain(v) for k, v in value.items()} if isinstance(value, (list, tuple)): return [_plain(v) for v in value] return value def _require(ok, message): if not ok: raise PretrainedError(message) def _read_allowed(path, roots, limit): path = Path(path).expanduser().resolve(strict=True) _require( any(path.is_relative_to(root) for root in roots), "Source is outside configured allowed roots", ) _require(path.is_file() and path.stat().st_size <= limit, "Source file missing or oversized") # Hash and deserialize the same bytes, not a second pathname lookup. with path.open("rb") as stream: data = stream.read(limit + 1) _require(len(data) <= limit, "Source file oversized") return path, data def _yaml_data(data): try: result = yaml.load(data, Loader=_DataLoader) _require(isinstance(result, dict), "Expected YAML mapping") return result except yaml.YAMLError as exc: raise PretrainedError("Unsupported or invalid configuration YAML") from exc def _actor_config(config): actor = dict(config["actor"]) distribution = dict(actor["distribution_cfg"]) # Old RSL config omitted this default; the tensor contract is checked too. distribution.setdefault("class_name", "GaussianDistribution") actor["distribution_cfg"] = distribution actor.setdefault("class_name", "MLPModel") actor.setdefault("cnn_cfg", None) return actor def validate_semantics(source_env, source_agent, target_env, target_agent): """Compare source and target against the repository's supported legacy47 contract.""" from src.tasks.velocity.config.go2.env_cfgs import unitree_go2_flat_env_cfg from src.tasks.velocity.config.go2.rl_cfg import unitree_go2_ppo_runner_cfg reference = _plain(asdict(unitree_go2_flat_env_cfg())) agent = _plain(asdict(unitree_go2_ppo_runner_cfg())) target_env, target_agent = _plain(target_env), _plain(target_agent) try: for label, env in (("source", source_env), ("target", target_env)): group = env["observations"]["actor"] names = list(group["terms"]) _require(names[:7] == BASE_TERMS, f"{label}: unknown base observation order") _require( names == BASE_TERMS if label == "source" else names in (BASE_TERMS, BASE_TERMS + ["forward_depth", "target_error"]), f"{label}: unsupported observation suffix", ) for name in BASE_TERMS: _require( group["terms"][name] == reference["observations"]["actor"]["terms"][name], f"{label}: incompatible observation {name}", ) for key in ( "concatenate_terms", "concatenate_dim", "history_length", "flatten_history_dim", ): _require( group[key] == reference["observations"]["actor"][key], f"{label}: incompatible observation {key}", ) _require( env["actions"] == reference["actions"], f"{label}: incompatible action transform" ) robot, ref_robot = ( env["scene"]["entities"]["robot"], reference["scene"]["entities"]["robot"], ) for key in ("spec_fn", "articulation", "sort_actuators", "collisions"): _require(robot[key] == ref_robot[key], f"{label}: incompatible robot {key}") for key in ("joint_pos", "joint_vel"): _require( robot["init_state"][key] == ref_robot["init_state"][key], f"{label}: incompatible default {key}", ) _require( env["decimation"] == 4 and env["sim"]["mujoco"]["timestep"] == 0.005, f"{label}: policy must run at 50 Hz", ) if len(target_env["observations"]["actor"]["terms"]) > 7: from src.tasks.obstacle_avoidance.env_cfg import unitree_go2_obstacle_env_cfg suffix = _plain(asdict(unitree_go2_obstacle_env_cfg()))["observations"]["actor"][ "terms" ] terms = target_env["observations"]["actor"]["terms"] depth = dict(terms["forward_depth"]) _require( set(depth["params"]) == {"max_distance"} and 0 < depth["params"]["max_distance"] < float("inf"), "Invalid ray normalization", ) depth["params"] = suffix["forward_depth"]["params"] _require( depth == suffix["forward_depth"] and terms["target_error"] == suffix["target_error"], "Unknown ray/goal observation semantics", ) for label, cfg in (("source", source_agent), ("target", target_agent)): _require( _actor_config(cfg) == _actor_config(agent), f"{label}: incompatible actor architecture/activation/std", ) _require( cfg["obs_groups"]["actor"] == ["actor"] and cfg["clip_actions"] is None, f"{label}: incompatible actor groups/action clipping", ) except (KeyError, TypeError) as exc: raise PretrainedError(f"Missing or invalid semantic configuration: {exc}") from exc def make_reference_actor(dim=47): from rsl_rl.models import MLPModel from tensordict import TensorDict return MLPModel( TensorDict({"actor": torch.zeros(1, dim)}, batch_size=[1]), {"actor": ["actor"]}, "actor", 12, hidden_dims=(512, 256, 128), activation="elu", obs_normalization=True, distribution_cfg={ "class_name": "GaussianDistribution", "init_std": 1.0, "std_type": "scalar", }, ) def validate_actor_state(state, dim=47): expected = make_reference_actor(dim).state_dict() _require( isinstance(state, dict) and state.keys() == expected.keys(), "Unexpected actor state keys" ) for key, tensor in state.items(): _require( isinstance(tensor, torch.Tensor) and tensor.shape == expected[key].shape and tensor.dtype == expected[key].dtype, f"Incompatible actor tensor: {key}", ) _require(bool(torch.isfinite(tensor).all()), f"Nonfinite actor tensor: {key}") _require(state["obs_normalizer.count"].item() >= 0, "Negative normalizer count") _require( bool((state["obs_normalizer._var"] >= 0).all()) and bool((state["obs_normalizer._std"] >= 0).all()), "Negative normalization variance/std", ) _require( torch.allclose( state["obs_normalizer._std"].square(), state["obs_normalizer._var"], atol=1e-5, rtol=1e-5, ), "Inconsistent normalization variance/std", ) _require(bool((state["distribution.std_param"] > 0).all()), "Nonpositive exploration std") def comparison_observations(): """Deterministic random + physically plausible standing/walking probe vectors.""" g = torch.Generator().manual_seed(20260825) random = torch.randn(32, 47, generator=g) physical = torch.zeros(16, 47) physical[:, 5] = -1 # body-frame projected gravity physical[:, 6] = torch.linspace(0, 1, 16) phase = torch.linspace(0, 2 * torch.pi, 16) physical[:, 9], physical[:, 10] = phase.sin(), phase.cos() physical[0, 9:11] = 0 # standing phase is masked return torch.cat((random, physical)) def verify_onnx(onnx_bytes, actor): import numpy as np import onnx import onnxruntime as ort graph = onnx.load_model_from_string(onnx_bytes) _require( all(t.data_location != onnx.TensorProto.EXTERNAL for t in graph.graph.initializer), "External ONNX tensors are not allowed", ) options = ort.SessionOptions() options.intra_op_num_threads = 1 options.inter_op_num_threads = 1 session = ort.InferenceSession( onnx_bytes, sess_options=options, providers=["CPUExecutionProvider"] ) inputs, outputs = session.get_inputs(), session.get_outputs() _require( len(inputs) == len(outputs) == 1 and inputs[0].shape == [1, 47] and outputs[0].shape == [1, 12] and inputs[0].type == outputs[0].type == "tensor(float)", "ONNX must be float32 [1,47] -> [1,12]", ) metadata = session.get_modelmeta().custom_metadata_map _require( metadata.get("joint_names", "").split(",") == JOINTS and metadata.get("observation_names", "").split(",") == BASE_TERMS and metadata.get("command_names") == "twist", "ONNX joint/observation/command semantics mismatch", ) for key, expected in { "action_scale": [0.25], "joint_stiffness": [20, 20, 40] * 4, "joint_damping": [1, 1, 2] * 4, "default_joint_pos": [-0.1, 0.9, -1.8, 0.1, 0.9, -1.8] * 2, }.items(): try: actual = [float(v) for v in metadata[key].split(",")] _require( len(actual) == len(expected) and np.allclose(actual, expected, rtol=0, atol=1e-6), f"ONNX {key} mismatch", ) except (KeyError, ValueError) as exc: raise PretrainedError(f"Invalid ONNX {key}") from exc actor.eval() obs = comparison_observations() with torch.inference_mode(): expected = actor.mlp(actor.obs_normalizer(obs)).numpy() actual = np.concatenate( [session.run(None, {inputs[0].name: row[None].numpy()})[0] for row in obs] ) _require( np.isfinite(actual).all() and np.allclose(actual, expected, atol=2e-5, rtol=2e-5), "ONNX does not match the selected checkpoint actor", ) return { "provider": "CPUExecutionProvider", "probe_count": len(obs), "max_abs_error": float(np.max(np.abs(actual - expected))), "atol": 2e-5, "rtol": 2e-5, } @dataclass class ValidatedSource: actor_state: dict manifest: dict def read_pretrained_source(checkpoint, *, allowed_roots, target_env, target_agent, onnx_path=None): """Read an explicit .pt plus paired ONNX/config. Never guess checkpoint by mtime. Future services should resolve opaque source IDs to these server-owned paths. An ONNX selection without an explicit corresponding .pt is not a training source. """ roots = [Path(p).expanduser().resolve(strict=True) for p in allowed_roots] _require( bool(roots) and all(p.is_dir() for p in roots), "Configure at least one local allowed source root", ) checkpoint, checkpoint_bytes = _read_allowed(checkpoint, roots, 256 * 1024 * 1024) _require( checkpoint.suffix == ".pt", "Warm-start requires a corresponding .pt checkpoint; ONNX cannot resume PPO", ) files = {"checkpoint": (checkpoint, checkpoint_bytes)} for name, path, limit in ( ("onnx", onnx_path or checkpoint.parent / "policy.onnx", 64 * 1024 * 1024), ("env", checkpoint.parent / "params/env.yaml", 2 * 1024 * 1024), ("agent", checkpoint.parent / "params/agent.yaml", 128 * 1024), ): files[name] = _read_allowed(path, roots, limit) source_env, source_agent = _yaml_data(files["env"][1]), _yaml_data(files["agent"][1]) validate_semantics(source_env, source_agent, target_env, target_agent) try: state = torch.load(io.BytesIO(checkpoint_bytes), map_location="cpu", weights_only=True) _require( isinstance(state, dict) and isinstance(state.get("iter"), int), "Invalid checkpoint iteration", ) actor_state = state["actor_state_dict"] validate_actor_state(actor_state) except (KeyError, RuntimeError) as exc: raise PretrainedError("Unsupported checkpoint") from exc actor = make_reference_actor() actor.load_state_dict(actor_state, strict=True) identity = verify_onnx(files["onnx"][1], actor) artifacts = { name: {"name": path.name, "sha256": hashlib.sha256(data).hexdigest(), "bytes": len(data)} for name, (path, data) in files.items() } manifest = { "schema_version": 1, "mode": "pretrained-warm-start", "contract": "go2-legacy47-v1", "artifacts": artifacts, "source_iteration": state["iter"], "source_actor_dim": 47, "source_normalizer_count": actor_state["obs_normalizer.count"].item(), "verification_versions": { name: version(name) for name in ("torch", "rsl-rl-lib", "mjlab", "onnxruntime", "PyYAML") }, "normalization": NORMALIZATION_POLICY, "onnx_identity": identity, "critic": "fresh-target-initialization", "optimizer": "fresh", "iteration": 0, "base_observation_terms": BASE_TERMS, "joint_names": JOINTS, } manifest["source_id"] = hashlib.sha256( json.dumps(artifacts, sort_keys=True).encode() ).hexdigest() return ValidatedSource(actor_state, manifest) def warm_start_actor(actor, source): """Strictly transfer a validated legacy actor into a fresh 47/81/97 actor.""" dim = actor.obs_dim _require(dim in (47, 81, 97), "Only legacy47 and ray/goal 81/97 actors are supported") reference = make_reference_actor(dim) _require( type(actor) is type(reference) and repr(actor.mlp) == repr(reference.mlp) and type(actor.distribution) is type(reference.distribution) and actor.distribution.std_type == "scalar" and actor.obs_normalization and type(actor.obs_normalizer) is type(reference.obs_normalizer) and actor.obs_normalizer.eps == reference.obs_normalizer.eps and actor.obs_normalizer.until is None and list(actor.obs_groups) == ["actor"], "Target actor runtime architecture mismatch", ) validate_actor_state(actor.state_dict(), dim) validate_actor_state(source.actor_state) migrated = {} for key, tensor in source.actor_state.items(): if key == "mlp.0.weight": value = tensor.new_zeros((512, dim)) value[:, :47] = tensor elif key in ("obs_normalizer._mean", "obs_normalizer._var", "obs_normalizer._std"): value = tensor.new_full((1, dim), 0 if key.endswith("_mean") else 1) value[:, :47] = tensor else: value = tensor.clone() migrated[key] = value actor.load_state_dict(migrated, strict=True) return { **source.manifest, "target_actor_dim": dim, "copied_tensors": list(source.actor_state), "zero_initialized_input_columns": [47, dim], "new_feature_statistics": {"mean": 0, "var": 1, "std": 1}, } def validate_runtime_contract(env): """Confirm compiled joint order/PD/defaults, not only configuration names.""" from mjlab.rl.exporter_utils import get_base_metadata metadata = get_base_metadata(env, "pretrained-validation") _require(metadata["joint_names"] == JOINTS, "Compiled robot joint order mismatch") _require( list(metadata["observation_names"]) in (BASE_TERMS, BASE_TERMS + ["forward_depth", "target_error"]), "Compiled observation order mismatch", ) _require(metadata["command_names"] == ["twist"], "Compiled command mismatch") for key, expected in { "action_scale": [0.25], "joint_stiffness": [20, 20, 40] * 4, "joint_damping": [1, 1, 2] * 4, "default_joint_pos": [-0.1, 0.9, -1.8, 0.1, 0.9, -1.8] * 2, }.items(): actual = torch.as_tensor(metadata[key], dtype=torch.float64).reshape(-1) target = torch.tensor(expected, dtype=torch.float64) _require( actual.shape == target.shape and torch.allclose(actual, target, atol=1e-6, rtol=0), f"Compiled {key} mismatch", ) return metadata def public_initialization_metadata(initialization): """Export provenance without local paths, including original single-file identity.""" artifacts = initialization["artifacts"] result = { "sourceId": initialization["source_id"], "checkpointSha256": artifacts["checkpoint"]["sha256"], "normalization": initialization["normalization"], "sourceIteration": initialization["source_iteration"], "initialIteration": 0, } if "sourceFormat" in initialization: result.update( sourceFormat=initialization["sourceFormat"], uploadSha256=artifacts["upload"]["sha256"], templateId=initialization["contract"], derivedFields=initialization["derived_fields"], templateConfirmation=initialization["template_confirmation"], ) else: result["onnxSha256"] = artifacts["onnx"]["sha256"] return result def initialize_runner(runner, source): _require( runner.current_learning_iteration == 0 and not runner.alg.optimizer.state, "Warm-start requires a fresh runner, never a resumed trial", ) # Do not load critic/optimizer/iteration from the source, even if shapes match. return warm_start_actor(runner.alg.actor, source)