feat(training): release V0.9.1 避障训练与基础策略迁移
This commit is contained in:
@@ -0,0 +1,502 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user