"""Single-file Go2 legacy47 import. No adjacent files or arbitrary ONNX conversion.""" import hashlib import io import json from dataclasses import asdict from pathlib import Path import torch from pretrained import ( BASE_TERMS, JOINTS, NORMALIZATION_POLICY, PretrainedError, ValidatedSource, _require, make_reference_actor, validate_actor_state, validate_semantics, verify_onnx, ) TEMPLATE = "go2-legacy47-v1" SYNTHETIC_COUNT = 1_000_000 UPLOAD_LIMITS = {"pt": 256 * 1024**2, "onnx": 64 * 1024**2} SEMANTICS = { "joint_names": JOINTS, "observation_names": BASE_TERMS, "command_names": ["twist"], "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, } def validate_metadata(metadata, *, required=False): """Check recognized embedded semantics; absence remains a user template assumption.""" import numpy as np _require(isinstance(metadata, dict), "模型metadata必须是对象") checked = [] for key, expected in SEMANTICS.items(): if key not in metadata: _require(not required, f"ONNX缺少语义metadata: {key}") continue actual = metadata[key] if isinstance(actual, str): actual = actual.split(",") if isinstance(expected[0], str): _require(actual == expected, f"模型metadata冲突: {key}") else: try: actual = np.asarray(actual, dtype=np.float64) _require( actual.shape == (len(expected),) and np.isfinite(actual).all() and np.allclose(actual, expected, atol=1e-6, rtol=0), f"模型metadata冲突: {key}", ) except (TypeError, ValueError) as error: raise PretrainedError(f"模型metadata无效: {key}") from error checked.append(key) for key in ("contract", "templateId"): if key in metadata: _require(metadata[key] == TEMPLATE, "模型模板metadata冲突") checked.append(key) return checked def _checkpoint(data): try: checkpoint = torch.load(io.BytesIO(data), map_location="cpu", weights_only=True) except Exception as error: raise PretrainedError( "不支持或不安全的.pt；仅支持weights_only Go2 legacy47 actor checkpoint" ) from error _require( isinstance(checkpoint, dict) and "actor_state_dict" in checkpoint, "请选择含actor_state_dict的Go2 legacy47 .pt，不接受ZIP工程或任意模型", ) state = checkpoint["actor_state_dict"] if isinstance(state, dict) and isinstance(state.get("mlp.0.weight"), torch.Tensor): _require( tuple(state["mlp.0.weight"].shape) == (512, 47), "当前仅支持Go2 legacy47输入；81/97 checkpoint请使用原trial续训，而不是基础策略上传", ) validate_actor_state(state) iteration = checkpoint.get("iter") _require( iteration is None or (type(iteration) is int and iteration >= 0), "无效checkpoint iteration" ) checked = validate_metadata(checkpoint) for key in ("metadata", "infos"): if key in checkpoint: checked.extend(validate_metadata(checkpoint[key])) nested = checkpoint[key].get("metadata") if nested is not None: checked.extend(validate_metadata(nested)) return state, iteration, sorted(set(checked)) def _onnx(data): import onnx from onnx import helper, numpy_helper model = onnx.load_model_from_string(data) graph = model.graph _require( not model.functions and not model.training_info and len(model.opset_import) == 1 and model.opset_import[0].domain == "" and model.opset_import[0].version in (17, 18), "仅支持标准opset17/18受限MLP ONNX", ) _require( not graph.sparse_initializer and not graph.quantization_annotation, "不支持稀疏或量化ONNX" ) _require(len(graph.input) == len(graph.output) == 1, "ONNX必须单输入单输出") for value, shape in ((graph.input[0], [1, 47]), (graph.output[0], [1, 12])): tensor = value.type.tensor_type _require( tensor.elem_type == onnx.TensorProto.FLOAT and [d.dim_value for d in tensor.shape.dim] == shape and all(not d.dim_param for d in tensor.shape.dim), "ONNX仅支持float32 [1,47] -> [1,12]；其他输入请使用对应训练器", ) expected_shapes = {"obs_normalizer._mean": (1, 47), "onnx::Div_24": (1, 47)} for i, shape in zip((0, 2, 4, 6), ((512, 47), (256, 512), (128, 256), (12, 128)), strict=True): expected_shapes[f"mlp.{i}.weight"] = shape expected_shapes[f"mlp.{i}.bias"] = (shape[0],) _require( len(graph.initializer) == len(expected_shapes) and {t.name for t in graph.initializer} == set(expected_shapes), "ONNX initializer不符合受支持MLP", ) tensors = {} for tensor in graph.initializer: _require( tensor.data_location == onnx.TensorProto.DEFAULT and not tensor.external_data, "拒绝ONNX external data；必须是单个自包含文件", ) _require( tensor.data_type == onnx.TensorProto.FLOAT and tuple(tensor.dims) == expected_shapes[tensor.name], "ONNX tensor类型/shape不支持", ) tensors[tensor.name] = torch.from_numpy(numpy_helper.to_array(tensor).copy()) _require(bool(torch.isfinite(tensors[tensor.name]).all()), "ONNX tensor含非有限值") # Exact dataflow, not merely op/tensor names: no branch, reorder, alias or extra op. expected_ops = ["Sub", "Div", "Gemm", "Elu", "Gemm", "Elu", "Gemm", "Elu", "Gemm"] _require( [n.op_type for n in graph.node] == expected_ops, "ONNX必须是Sub/Div及4层Gemm+3层ELU精确链路" ) previous = graph.input[0].name seen = set(tensors) | {previous} _require(len(seen) == len(tensors) + 1, "ONNX输入与initializer重名") layer = 0 for node in graph.node: _require(node.domain == "" and not node.overload, "拒绝ONNX custom op") attributes = {a.name: helper.get_attribute_value(a) for a in node.attribute} _require(len(attributes) == len(node.attribute), "重复ONNX属性") if node.op_type in ("Sub", "Div"): inputs = [previous, "obs_normalizer._mean" if node.op_type == "Sub" else "onnx::Div_24"] _require(not attributes, "不支持normalizer算子属性") elif node.op_type == "Gemm": inputs = [previous, f"mlp.{layer}.weight", f"mlp.{layer}.bias"] layer += 2 _require( set(attributes) <= {"alpha", "beta", "transA", "transB"} and attributes.get("alpha", 1.0) == 1.0 and attributes.get("beta", 1.0) == 1.0 and attributes.get("transA", 0) == 0 and attributes.get("transB", 0) == 1, "不支持Gemm缩放/转置属性", ) else: inputs = [previous] _require( set(attributes) <= {"alpha"} and attributes.get("alpha", 1.0) == 1.0, "仅支持ELU alpha=1", ) _require( list(node.input) == inputs and len(node.output) == 1 and node.output[0] and node.output[0] not in seen, "ONNX实际连边/输出不符合受支持MLP", ) previous = node.output[0] seen.add(previous) _require(previous == graph.output[0].name, "ONNX输出必须是最后Gemm结果") metadata = {p.key: p.value for p in model.metadata_props} _require(len(metadata) == len(model.metadata_props), "重复ONNX metadata") checked = validate_metadata(metadata, required=True) onnx.checker.check_model(model, full_check=True) actor = make_reference_actor() _require( actor.obs_normalizer.eps == 0.01 and bool((actor.state_dict()["distribution.std_param"] == 1).all()), "目标训练默认normalizer/exploration已变化，需要新模板", ) state = actor.state_dict() for key in state: if key in tensors: state[key] = tensors[key] std = tensors["onnx::Div_24"] - 0.01 _require(bool((std > 0).all()), "ONNX denominator必须大于模板epsilon=.01") state["obs_normalizer._std"] = std state["obs_normalizer._var"] = std.square() state["obs_normalizer.count"].fill_(SYNTHETIC_COUNT) validate_actor_state(state) actor.load_state_dict(state) identity = verify_onnx(data, actor) return state, checked, identity def source_identity(fmt, digest): return hashlib.sha256(f"{TEMPLATE}:{fmt}:{digest}".encode()).hexdigest() def import_upload(path, fmt, template, directory): """Run only inside the resource-limited validator process.""" _require(template == TEMPLATE, "必须明确确认go2-legacy47-v1模板") _require(fmt in UPLOAD_LIMITS, "仅支持单个.pt或.onnx，不支持ZIP") data = Path(path).read_bytes() _require(0 < len(data) <= UPLOAD_LIMITS[fmt], "上传文件为空或过大") if fmt == "pt": state, iteration, checked = _checkpoint(data) identity = None else: state, checked, identity = _onnx(data) iteration = None from pretrained import comparison_observations actor = make_reference_actor().eval() actor.load_state_dict(state) with torch.inference_mode(): output = actor.mlp(actor.obs_normalizer(comparison_observations())) _require(bool(torch.isfinite(output).all()), "actor在随机/物理probe上产生非有限动作") directory = Path(directory) actor_path = directory / "actor.pt" torch.save({"actor_state_dict": state}, actor_path) digest = hashlib.sha256(data).hexdigest() manifest = { "schema_version": 1, "mode": "pretrained-warm-start", "contract": TEMPLATE, "sourceFormat": fmt, "source_id": source_identity(fmt, digest), "source_iteration": iteration, "source_actor_dim": 47, "source_normalizer_count": state["obs_normalizer.count"].item(), "normalization": NORMALIZATION_POLICY if fmt == "pt" else "synthetic-count/unit-new-features", "template_confirmation": { "id": TEMPLATE, "confirmed_by": "user", "assumptions": ( "Go2 legacy47 observation physics, ordering, 50Hz and action semantics; " "not verified source env.yaml" ), }, "verified_facts": { "actor_tensor_shapes": "47-512-256-128-12/float32", "metadata_fields": checked, "finite_output_probes": 48, "activation": "graph-verified-ELU" if fmt == "onnx" else "user-template-assumed-ELU", }, "derived_fields": {} if fmt == "pt" else { "normalizer_count": {"policy": "synthetic", "value": SYNTHETIC_COUNT}, "normalizer_std": "denominator - 0.01", "normalizer_var": "std squared", "epsilon": 0.01, "exploration_std": {"policy": "fresh-target-default", "value": 1.0}, }, "onnx_identity": identity, "critic": "fresh-target-initialization", "optimizer": "fresh", "iteration": 0, "base_observation_terms": BASE_TERMS, "joint_names": JOINTS, "artifacts": { "upload": {"name": f"upload.{fmt}", "sha256": digest, "bytes": len(data)}, "checkpoint": { "name": "actor.pt", "sha256": hashlib.sha256(actor_path.read_bytes()).hexdigest(), "bytes": actor_path.stat().st_size, "origin": "service-derived-actor-only", }, }, } return manifest def read_uploaded_source(checkpoint, *, allowed_roots, manifest_path, target_env, target_agent): """Load only the service-derived actor and bound provenance, never adjacent sidecars.""" from pretrained_sources import regular_bytes roots = [Path(p).resolve(strict=True) for p in allowed_roots] checkpoint, manifest_path = Path(checkpoint), Path(manifest_path) root = next( (r for r in roots if checkpoint.is_relative_to(r) and manifest_path.is_relative_to(r)), None ) _require(root is not None, "上传artifact不在受控根内") manifest = json.loads(regular_bytes(manifest_path, root, 64 * 1024)) _require( manifest.get("contract") == TEMPLATE and manifest.get("sourceFormat") in UPLOAD_LIMITS, "无效上传manifest模板/格式", ) artifacts = manifest["artifacts"] _require( manifest["source_id"] == source_identity(manifest["sourceFormat"], artifacts["upload"]["sha256"]), "上传原始SHA身份不匹配", ) data = regular_bytes(checkpoint, root, UPLOAD_LIMITS["pt"]) _require( hashlib.sha256(data).hexdigest() == artifacts["checkpoint"]["sha256"], "上传actor SHA不匹配" ) from pretrained import _plain 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 validate_semantics( _plain(asdict(unitree_go2_flat_env_cfg())), _plain(asdict(unitree_go2_ppo_runner_cfg())), target_env, target_agent, ) state, _, _ = _checkpoint(data) return ValidatedSource(state, manifest)