503 lines
19 KiB
Python
503 lines
19 KiB
Python
"""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)
|