feat(training): release V0.9.1 避障训练与基础策略迁移

training_server/rl/requirements.txt

@@ -7,3 +7,7 @@ mujoco-warp==3.5.0
 warp-lang==1.15.0
 # RSL-RL 5.0.1 仍向 wandb.Settings 传递 start_method，0.29 已删除该字段。
 wandb==0.28.2
+# 基础策略迁移依赖已验证的 actor/normalizer/checkpoint 契约。
+rsl-rl-lib==5.0.1
+# 基础策略身份校验只使用 CPU ONNX Runtime。
+onnxruntime==1.29.0

training_server/rl/scripts/evaluate.py

@@ -32,7 +32,6 @@ from mjlab.tasks.registry import list_tasks, load_env_cfg, load_rl_cfg, load_run
 from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
 from mjlab.utils.torch import configure_torch_backends
 from torch.utils.tensorboard import SummaryWriter
-
 from tuning.schema import apply_reward_configuration, validate_configuration
 
 SCENARIOS = (
@@ -61,6 +60,7 @@ class EvaluateConfig:
   checkpoint: str
   output: str
   reward_config: str | None = None
+  task_config: str | None = None
   num_envs: int = 256
   steps_per_seed: int = 1000
   seeds: tuple[int, ...] = field(default_factory=lambda: (101, 202, 303))
@@ -163,31 +163,41 @@ def run_evaluation(task_id: str, cfg: EvaluateConfig) -> dict:
   checkpoint = Path(cfg.checkpoint).expanduser().resolve(strict=True)
   output = Path(cfg.output).expanduser().resolve()
   output.parent.mkdir(parents=True, exist_ok=True)
-  per_seed = [_evaluate_seed(task_id, cfg, seed) for seed in cfg.seeds]
-  metrics = {
-    key: fmean(seed_metrics[key] for seed_metrics in per_seed)
-    for key in (*METRIC_NAMES, "fall_rate")
-  }
-  deviations = {
-    key: pstdev(seed_metrics[key] for seed_metrics in per_seed)
-    for key in (*METRIC_NAMES, "fall_rate")
-  }
-  result = {
-    "protocolVersion": 1,
-    "taskId": task_id,
-    "checkpoint": checkpoint.name,
-    "checkpointSha256": _sha256(checkpoint),
-    "seeds": list(cfg.seeds),
-    "numEnvs": cfg.num_envs,
-    "stepsPerSeed": cfg.steps_per_seed,
-    "scenarios": [list(value) for value in SCENARIOS],
-    "metrics": metrics,
-    "metricStd": deviations,
-    "seedMetrics": [
-      {"seed": seed, "metrics": values}
-      for seed, values in zip(cfg.seeds, per_seed, strict=True)
-    ],
-  }
+  if task_id == "Unitree-Go2-ObstacleAvoidance":
+    from scripts.evaluate_obstacle import run_obstacle_evaluation
+    checkpoint_hash = _sha256(checkpoint)
+    result = run_obstacle_evaluation(task_id, cfg)
+    if _sha256(checkpoint) != checkpoint_hash:
+      raise ValueError("Checkpoint changed during evaluation")
+    result.update({"taskId": task_id, "checkpoint": checkpoint.name,
+                   "checkpointSha256": checkpoint_hash})
+    metrics = result["metrics"]
+  else:
+    per_seed = [_evaluate_seed(task_id, cfg, seed) for seed in cfg.seeds]
+    metrics = {
+      key: fmean(seed_metrics[key] for seed_metrics in per_seed)
+      for key in (*METRIC_NAMES, "fall_rate")
+    }
+    deviations = {
+      key: pstdev(seed_metrics[key] for seed_metrics in per_seed)
+      for key in (*METRIC_NAMES, "fall_rate")
+    }
+    result = {
+      "protocolVersion": 1,
+      "taskId": task_id,
+      "checkpoint": checkpoint.name,
+      "checkpointSha256": _sha256(checkpoint),
+      "seeds": list(cfg.seeds),
+      "numEnvs": cfg.num_envs,
+      "stepsPerSeed": cfg.steps_per_seed,
+      "scenarios": [list(value) for value in SCENARIOS],
+      "metrics": metrics,
+      "metricStd": deviations,
+      "seedMetrics": [
+        {"seed": seed, "metrics": values}
+        for seed, values in zip(cfg.seeds, per_seed, strict=True)
+      ],
+    }
   output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
   writer = SummaryWriter(log_dir=str(output.parent / "evaluation-events"))
   try:

training_server/rl/scripts/evaluate_obstacle.py

@@ -0,0 +1,241 @@
+"""Actual checkpoint inference, capturing first-terminal physics before mjlab resets."""
+
+import hashlib
+import json
+import shutil
+import subprocess
+import sys
+import tempfile
+from dataclasses import asdict
+from pathlib import Path
+from statistics import fmean
+from types import SimpleNamespace
+
+# Also executable as a fresh-interpreter seed worker (never fork a CUDA context).
+for source in (Path(__file__).resolve().parents[1], Path(__file__).resolve().parents[2]):
+    if str(source) not in sys.path:
+        sys.path.insert(0, str(source))
+
+import torch
+from mjlab.envs import ManagerBasedRlEnv
+from mjlab.rl import MjlabOnPolicyRunner, RslRlVecEnvWrapper
+from mjlab.sensor import ContactMatch, ContactSensorCfg
+from mjlab.tasks.registry import load_env_cfg, load_rl_cfg, load_runner_cls
+from scripts.train import _load_reward_config, _load_task_config
+from src.tasks.obstacle_avoidance.env_cfg import apply_obstacle_configuration
+from task_config import OBSTACLE_TASK
+from tuning.obstacle_scoring import METRICS, STEPS, protocol, score_trajectory
+from tuning.schema import apply_reward_configuration
+
+
+class FirstEpisodeRecorder:
+    """The termination hook runs post-physics, before _reset_idx (including first terminal)."""
+
+    def __init__(self, env, layout):
+        self.env = env
+        self.layout = layout
+        self.samples = [[] for _ in range(env.num_envs)]
+        self.previous = torch.zeros((env.num_envs, 12), device=env.device)
+        self.calls = 0
+        self.compute = env.termination_manager.compute
+        env.termination_manager.compute = self.capture
+
+    def capture(self):
+        env = self.env
+        terminal = self.compute()
+        robot = env.scene["robot"].data
+        # Like mjlab termination itself, derived pose is one physics substep old.
+        position = robot.root_link_pos_w
+        local_xy = position[:, :2] - env.scene.env_origins[:, :2]
+        xy = local_xy + torch.tensor(self.layout["spawn"][:2], device=env.device)
+        boxes = self.layout["boxes"][1:]  # Never include the support floor.
+        clearance = torch.full((env.num_envs,), 0.5, device=env.device)
+        if boxes:
+            centers = torch.tensor([b["pos"][:2] for b in boxes], device=env.device)
+            sizes = torch.tensor([b["size"][:2] for b in boxes], device=env.device)
+            clearance = (
+                ((xy[:, None, :] - centers).abs() - sizes).clamp(min=0).norm(dim=2).amin(dim=1)
+            )
+            clearance = (clearance - 0.3).clamp(min=0)  # Conservative body footprint radius.
+            forces = env.scene["evaluation_obstacles"].data.force_history
+            collision = forces.norm(dim=-1).flatten(1).amax(dim=1) > 1.0
+        else:
+            collision = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
+        nonfoot = env.scene["nonfoot_ground_touch"].data.force_history
+        collision |= nonfoot.norm(dim=-1).flatten(1).amax(dim=1) > 10.0
+        fall = (robot.projected_gravity_b[:, 2] > -0.3420201433) | (position[:, 2] < 0.12)
+        actions = env.action_manager.action
+        delta = (actions - self.previous).square().mean(dim=1)
+        self.previous.copy_(actions)
+        rays = env.scene["forward_scan"].data.distances
+        hits = (
+            ((rays >= 0) & (rays <= env.scene["forward_scan"].cfg.max_distance)).float().mean(dim=1)
+        )
+        distance = env.command_manager.get_term("twist").errors()[1]
+        values = torch.stack((distance, clearance, delta, hits, collision, fall, terminal), dim=1)
+        if not torch.isfinite(values).all():
+            raise ValueError("Nonfinite rollout measurements")
+        keys = ("distance", "clearance", "action_delta", "ray_hit", "collision", "fall", "terminal")
+        for samples, row in zip(self.samples, values.cpu().tolist(), strict=True):
+            if not samples or not samples[-1]["terminal"]:
+                samples.append(dict(zip(keys, row, strict=True)))
+        self.calls += 1
+        return terminal
+
+    def metrics(self, horizon=STEPS):
+        if self.calls != horizon:
+            raise ValueError("Incomplete rollout horizon")
+        metrics = [score_trajectory(samples, horizon) for samples in self.samples]
+        return {key: fmean(item[key] for item in metrics) for key in METRICS}
+
+
+def configure_seed(task_id, cfg, scenario, reward):
+    env_cfg = load_env_cfg(task_id, play=False)
+    env_cfg.seed = scenario["seed"]
+    env_cfg.scene.num_envs = cfg.num_envs
+    # Keep the benchmark's declared pair fixed; training itself randomizes every reset.
+    apply_obstacle_configuration(env_cfg, scenario["taskConfig"], randomize_navigation=False)
+    apply_reward_configuration(env_cfg, reward, task_id)
+    env_cfg.curriculum = {}
+    env_cfg.observations["actor"].enable_corruption = False
+    env_cfg.events.pop("push_robot", None)
+    # Primary geom names are literal compiled static terrain geoms, not a regex.
+    names = tuple(f"terrain_{i}" for i in range(1, len(scenario["terrain"]["boxes"])))
+    if names:
+        env_cfg.scene.sensors += (
+            ContactSensorCfg(
+                name="evaluation_obstacles",
+                primary=ContactMatch(mode="geom", pattern=names),
+                secondary=ContactMatch(mode="subtree", pattern="base_link", entity="robot"),
+                fields=("force",),
+                reduce="maxforce",
+                history_length=env_cfg.decimation,
+            ),
+        )
+    return env_cfg
+
+
+def evaluate_seed(task_id, cfg, scenario, reward):
+    torch.manual_seed(scenario["seed"])
+    env_cfg = configure_seed(task_id, cfg, scenario, reward)
+    agent_cfg = load_rl_cfg(task_id)
+    env = ManagerBasedRlEnv(env_cfg, device=cfg.device)
+    wrapped = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
+    try:
+        runner_cls = load_runner_cls(task_id) or MjlabOnPolicyRunner
+        runner = runner_cls(wrapped, asdict(agent_cfg), log_dir=None, device=wrapped.device)
+        # rsl_rl actor state_dict includes observation_normalizer buffers. Strict load
+        # rejects incompatible architecture/statistics; no fixture or zero-action fallback.
+        runner.load(
+            str(Path(cfg.checkpoint).resolve(strict=True)),
+            load_cfg={"actor": True},
+            strict=True,
+            map_location=str(wrapped.device),
+        )
+        policy = runner.get_inference_policy(device=str(wrapped.device))
+        obs, _ = env.reset(seed=scenario["seed"])
+        recorder = FirstEpisodeRecorder(env, scenario["terrain"])
+        with torch.inference_mode():
+            for _ in range(STEPS):
+                actions = policy(obs)
+                if not torch.isfinite(actions).all():
+                    raise ValueError("Nonfinite policy actions")
+                obs, _, _, _ = wrapped.step(actions)
+        return {
+            "seed": scenario["seed"],
+            "metrics": recorder.metrics(),
+            "episodes": cfg.num_envs,
+            "rolloutSteps": STEPS,
+        }
+    finally:
+        wrapped.close()
+
+
+def run_obstacle_evaluation(task_id, cfg):
+    from tuning.obstacle_scoring import SEEDS
+
+    if tuple(cfg.seeds) != SEEDS or cfg.steps_per_seed != STEPS:
+        raise ValueError("Obstacle evaluation seeds/horizon are fixed")
+    if not cfg.task_config or not cfg.reward_config:
+        raise ValueError("Obstacle evaluation requires task and parameter configurations")
+    raw = json.loads(Path(cfg.task_config).read_text())
+    custom = _load_task_config(task_id, cfg.task_config, raw["seed"])
+    reward = _load_reward_config(cfg.reward_config, None, OBSTACLE_TASK)
+    fixed = protocol(custom, cfg.num_envs)
+    seeds = evaluate_isolated_seeds(task_id, cfg, fixed, reward)
+    return {
+        "protocol": fixed,
+        "seedMetrics": seeds,
+        "metrics": {key: fmean(item["metrics"][key] for item in seeds) for key in METRICS},
+    }
+
+
+def file_hash(path):
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as stream:
+        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def evaluate_isolated_seeds(task_id, cfg, fixed, reward):
+    from tuning.obstacle_scoring import validate_evaluation
+
+    seeds = []
+    # Inherit the evaluation parent's process group: manager cancellation kills
+    # both parent and the active seed. subprocess.run kills/reaps on timeout too.
+    with tempfile.TemporaryDirectory(prefix="go2-obstacle-eval-") as directory:
+        root = Path(directory)
+        checkpoint = root / "checkpoint.pt"
+        source_hash = file_hash(cfg.checkpoint)
+        shutil.copyfile(cfg.checkpoint, checkpoint)
+        if file_hash(checkpoint) != source_hash:
+            raise ValueError("Checkpoint changed during snapshot")
+        for scenario in fixed["scenarios"]:
+            request = root / f"request-{scenario['seed']}.json"
+            output = root / f"result-{scenario['seed']}.json"
+            request.write_text(
+                json.dumps(
+                    {
+                        "taskId": task_id,
+                        "config": {**asdict(cfg), "checkpoint": str(checkpoint)},
+                        "scenario": scenario,
+                        "reward": reward,
+                        "checkpointSha256": source_hash,
+                    },
+                    allow_nan=False,
+                )
+            )
+            subprocess.run(
+                [sys.executable, str(Path(__file__).resolve()), str(request), str(output)],
+                check=True,
+                timeout=1200,
+                shell=False,
+            )
+            result = json.loads(output.read_text())
+            if (
+                result.pop("checkpointSha256", None) != source_hash
+                or file_hash(checkpoint) != source_hash
+            ):
+                raise ValueError("Seed checkpoint identity mismatch")
+            seeds.append(result)
+        candidate = {
+            "protocol": fixed,
+            "seedMetrics": seeds,
+            "metrics": {key: fmean(item["metrics"][key] for item in seeds) for key in METRICS},
+        }
+        validate_evaluation(candidate, fixed)
+    return seeds
+
+
+if __name__ == "__main__":
+    import src.tasks  # noqa: F401
+
+    request_path, output_path = map(Path, sys.argv[1:])
+    request = json.loads(request_path.read_text())
+    config = SimpleNamespace(**request["config"])
+    if file_hash(config.checkpoint) != request["checkpointSha256"]:
+        raise ValueError("Worker checkpoint identity mismatch")
+    result = evaluate_seed(request["taskId"], config, request["scenario"], request["reward"])
+    result["checkpointSha256"] = file_hash(config.checkpoint)
+    output_path.write_text(json.dumps(result, allow_nan=False))

training_server/rl/scripts/train.py

@@ -7,7 +7,7 @@ import sys
 from dataclasses import asdict, dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Literal, cast
+from typing import Literal
 
 # 训练器作为仓库内置子集直接从 scripts/ 启动，不要求额外执行 pip install -e。
 TRAINER_ROOT = Path(__file__).resolve().parents[1]
@@ -34,7 +34,7 @@ from mjlab.utils.gpu import select_gpus
 from mjlab.utils.os import dump_yaml, get_checkpoint_path
 from mjlab.utils.torch import configure_torch_backends
 from mjlab.utils.wrappers import VideoRecorder
-
+from task_config import OBSTACLE_TASK, deployment_metadata, validate_task_config
 from tuning.schema import apply_reward_configuration, validate_configuration
 
 
@@ -51,8 +51,14 @@ class TrainConfig:
   gpu_ids: list[int] | Literal["all"] | None = field(default_factory=lambda: [0])
   output_dir: str | None = None
   resume_checkpoint: str | None = None
+  pretrained_checkpoint: str | None = None
+  pretrained_upload_manifest: str | None = None
+  pretrained_onnx: str | None = None
+  pretrained_source_id: str | None = None
+  pretrained_allowed_roots: list[str] = field(default_factory=list)
   reward_config: str | None = None
   reward_config_json: str | None = None
+  task_config: str | None = None
 
   @staticmethod
   def from_task(task_id: str) -> "TrainConfig":
@@ -61,27 +67,96 @@ class TrainConfig:
     return TrainConfig(env=env_cfg, agent=agent_cfg)
 
 
-def _load_reward_config(path: str | None, inline: str | None) -> dict | None:
+def _load_reward_config(path: str | None, inline: str | None, task_id="Unitree-Go2-Flat") -> dict | None:
   if path is not None and inline is not None:
     raise ValueError("Use only one of reward_config and reward_config_json")
   if inline is not None:
     if len(inline.encode("utf-8")) > 64 * 1024:
       raise ValueError("Reward configuration is larger than 64 KiB")
-    return validate_configuration(json.loads(inline))
+    return validate_configuration(json.loads(inline), task_id)
   if path is None:
     return None
   source = Path(path).expanduser().resolve(strict=True)
   if source.stat().st_size > 64 * 1024:
     raise ValueError("Reward configuration is larger than 64 KiB")
   with source.open(encoding="utf-8") as stream:
-    return validate_configuration(json.load(stream))
+    return validate_configuration(json.load(stream), task_id)
+
+
+def _load_task_config(task_id: str, path: str | None, seed: int) -> dict | None:
+  if path is None:
+    return validate_task_config(task_id, {}, seed)
+  source = Path(path).expanduser().resolve(strict=True)
+  if source.stat().st_size > 128 * 1024:
+    raise ValueError("Task configuration is larger than 128 KiB")
+  payload = json.loads(source.read_text(encoding="utf-8"))
+  allowed = {"terrainPreset", "terrainParams", "sensorCfg", "seed", "customTerrainBoxes"}
+  if not isinstance(payload, dict) or payload.keys() - allowed:
+    raise ValueError("Unknown task configuration fields")
+  if (
+    isinstance(payload.get("seed"), bool)
+    or not isinstance(payload.get("seed"), int)
+    or payload.get("seed") != seed
+  ):
+    raise ValueError("Task configuration seed must equal the agent seed")
+  if payload.get("sensorCfg") is None:
+    payload.pop("sensorCfg", None)
+  return validate_task_config(task_id, payload, seed)
+
+
+def _configure_task_and_rewards(task_id: str, cfg: TrainConfig):
+  custom = _load_task_config(task_id, cfg.task_config, cfg.agent.seed)
+  if custom is not None:
+    from src.tasks.obstacle_avoidance.env_cfg import apply_obstacle_configuration
+    from src.tasks.obstacle_avoidance.terrain import apply_terrain_configuration
+    if task_id == OBSTACLE_TASK:
+      apply_obstacle_configuration(cfg.env, custom)
+    else:
+      apply_terrain_configuration(cfg.env, custom)
+  deployment = deployment_metadata(task_id, custom, cfg.agent.seed)
+  reward_config = _load_reward_config(cfg.reward_config, cfg.reward_config_json, task_id)
+  if reward_config is not None:
+    apply_reward_configuration(cfg.env, reward_config, task_id)
+    if task_id == OBSTACLE_TASK:
+      deployment["navigation"]["speed"] = reward_config["params"]["target_velocity"]
+      deployment["sensorCfg"]["avoidanceWeight"] = reward_config["weights"]["avoidance_weight"]
+
+  return deployment, reward_config
+
+
+def _load_pretrained(cfg: TrainConfig):
+  if cfg.pretrained_checkpoint is None:
+    if (cfg.pretrained_onnx is not None or cfg.pretrained_allowed_roots
+        or cfg.pretrained_source_id or cfg.pretrained_upload_manifest):
+      raise ValueError("Pretrained options require --pretrained-checkpoint (.pt), not ONNX alone")
+    return None
+  if cfg.resume_checkpoint is not None or cfg.agent.resume:
+    raise ValueError("Pretrained warm-start and resume are mutually exclusive")
+  from pretrained import read_pretrained_source
+
+  options = {"onnx_path": cfg.pretrained_onnx}
+  if cfg.pretrained_upload_manifest is not None:
+    if cfg.pretrained_onnx is not None:
+      raise ValueError("Uploaded actor cannot use adjacent ONNX sidecars")
+    from pretrained_upload import read_uploaded_source
+    read_pretrained_source = read_uploaded_source
+    options = {"manifest_path": cfg.pretrained_upload_manifest}
+  source = read_pretrained_source(
+    cfg.pretrained_checkpoint,
+    allowed_roots=cfg.pretrained_allowed_roots,
+    **options,
+    target_env=asdict(cfg.env),
+    target_agent=asdict(cfg.agent),
+  )
+  if cfg.pretrained_source_id is not None and source.manifest["source_id"] != cfg.pretrained_source_id:
+    raise ValueError("基础策略SHA身份变化，拒绝初始化")
+  return source
 
 
 def run_train(task_id: str, cfg: TrainConfig, log_dir: Path) -> None:
-  reward_config = _load_reward_config(cfg.reward_config, cfg.reward_config_json)
-  if reward_config is not None:
-    apply_reward_configuration(cfg.env, reward_config)
-
+  deployment, reward_config = _configure_task_and_rewards(task_id, cfg)
+  # Validate source identity/semantics before allocating a simulation or optimizer.
+  pretrained = _load_pretrained(cfg)
   cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
   if cuda_visible == "":
     device = "cpu"
@@ -135,6 +210,17 @@ def run_train(task_id: str, cfg: TrainConfig, log_dir: Path) -> None:
     cfg=cfg.env, device=device, render_mode="rgb_array" if cfg.video else None
   )
 
+  if pretrained is not None:
+    from pretrained import validate_runtime_contract
+
+    try:
+      validate_runtime_contract(env)
+    except Exception:
+      env.close()
+      raise
+
+  # The ONNX runner exports this exact snapshot beside and inside policy.onnx.
+  env.platform_deployment = deployment
   log_root_path = log_dir.parent  # Go up from specific run dir to experiment dir.
 
   resume_path: Path | None = None
@@ -171,9 +257,30 @@ def run_train(task_id: str, cfg: TrainConfig, log_dir: Path) -> None:
   runner = runner_cls(env, agent_cfg, str(log_dir), device, **runner_kwargs)
 
   runner.add_git_repo_to_log(__file__)
+  if pretrained is not None:
+    from pretrained import initialize_runner
+
+    initialization = initialize_runner(runner, pretrained)
+    env.unwrapped.platform_initialization = initialization
+    if rank == 0:
+      (log_dir / "initialization.json").write_text(
+        json.dumps(initialization, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
+      )
+    print(f"[INFO] Warm-start actor from source {initialization['source_id']}; fresh critic/optimizer, iteration=0")
   if resume_path is not None:
     print(f"[INFO]: Loading model checkpoint from: {resume_path}")
     runner.load(str(resume_path), map_location=device)
+    # Promotion copies the origin manifest, never the original actor weights.
+    origin_path = resume_path.parent / "initialization.json"
+    if origin_path.is_file():
+      if origin_path.stat().st_size > 64 * 1024:
+        raise ValueError("Initialization provenance exceeds 64 KiB")
+      origin = json.loads(origin_path.read_text(encoding="utf-8"))
+      if not isinstance(origin, dict) or origin.get("mode") != "pretrained-warm-start":
+        raise ValueError("Invalid initialization provenance")
+      env.unwrapped.platform_initialization = origin
+      if rank == 0:
+        (log_dir / "initialization.json").write_text(json.dumps(origin, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
     if explicit_resume:
       # RSL-RL stores the last completed zero-based iteration and otherwise
       # repeats it after load. Explicit tuning promotion uses an absolute target.
@@ -260,7 +367,7 @@ def main():
   # Parse first argument to choose the task.
   # Import tasks to populate the registry.
   import mjlab.tasks  # noqa: F401
-  import src.tasks
+  import src.tasks  # noqa: F401
 
   all_tasks = list_tasks()
   chosen_task, remaining_args = tyro.cli(

training_server/rl/scripts/validate_pretrained.py

@@ -0,0 +1,48 @@
+"""Service-owned CPU validator; JSON stdin is never accepted directly from HTTP."""
+
+import json
+import sys
+from dataclasses import asdict
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+for root in (ROOT, ROOT.parent):
+    sys.path.insert(0, str(root))
+
+
+def main():
+    from pretrained import read_pretrained_source
+    from task_config import OBSTACLE_TASK
+    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
+
+    payload = json.loads(sys.stdin.read(128 * 1024))
+    task_id = payload["taskId"]
+    if task_id == "Unitree-Go2-Flat":
+        cfg = unitree_go2_flat_env_cfg()
+    elif task_id == OBSTACLE_TASK:
+        from src.tasks.obstacle_avoidance.env_cfg import unitree_go2_obstacle_env_cfg, apply_obstacle_configuration
+        cfg = unitree_go2_obstacle_env_cfg()
+        if payload.get("taskConfig") is not None:
+            apply_obstacle_configuration(cfg, payload["taskConfig"])
+    else:
+        raise ValueError("基础策略不兼容该任务；支持Flat47与Obstacle81/97，不支持Rough")
+    directory = Path(payload["directory"])
+    options = {}
+    if payload.get("uploaded"):
+        from pretrained_upload import read_uploaded_source
+        read_pretrained_source = read_uploaded_source
+        options["manifest_path"] = directory / "upload.json"
+    source = read_pretrained_source(
+        directory / payload["checkpoint"], allowed_roots=[directory], target_env=asdict(cfg),
+        target_agent=asdict(unitree_go2_ppo_runner_cfg()), **options,
+    )
+    print(json.dumps(source.manifest))
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as error:
+        print(f"基础策略不兼容或缺少配套文件/依赖：{error}", file=sys.stderr)
+        sys.exit(1)

training_server/rl/scripts/validate_upload.py

@@ -0,0 +1,43 @@
+"""Resource-limited CPU single-file decoder; no network or neighboring model lookup."""
+
+import contextlib
+import json
+import os
+import resource
+import sys
+from pathlib import Path
+
+# Apply limits before importing tensor/protobuf runtimes. Thread env also set by parent.
+resource.setrlimit(resource.RLIMIT_CPU, (40, 40))
+resource.setrlimit(resource.RLIMIT_AS, (8 * 1024**3, 8 * 1024**3))
+resource.setrlimit(resource.RLIMIT_FSIZE, (16 * 1024**2, 16 * 1024**2))
+resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64))
+resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+ROOT = Path(__file__).resolve().parents[1]
+for root in (ROOT, ROOT.parent):
+    sys.path.insert(0, str(root))
+
+if __name__ == "__main__":
+    try:
+        payload = json.loads(sys.stdin.read(4096))
+        with contextlib.redirect_stdout(sys.stderr):
+            import torch
+            from pretrained_upload import import_upload
+
+            torch.set_num_threads(1)
+            manifest = import_upload(
+                payload["path"], payload["format"], payload["template"], payload["directory"]
+            )
+        print(json.dumps(manifest))
+    except Exception as error:
+        # Do not echo untrusted pickle/tensor names or absolute local paths.
+        from pretrained import PretrainedError
+
+        message = (
+            str(error)
+            if isinstance(error, PretrainedError)
+            else "文件解析失败；请提供受支持的自包含Go2 legacy47 actor"
+        )
+        print(json.dumps({"error": message[:500]}))
+        sys.exit(1)

training_server/rl/src/tasks/obstacle_avoidance/__init__.py

@@ -0,0 +1,17 @@
+from mjlab.tasks.registry import register_mjlab_task
+from task_config import OBSTACLE_TASK
+
+from src.tasks.velocity.config.go2.rl_cfg import unitree_go2_ppo_runner_cfg
+from src.tasks.velocity.rl import VelocityOnPolicyRunner
+
+from .env_cfg import unitree_go2_obstacle_env_cfg
+
+_rl_cfg = unitree_go2_ppo_runner_cfg()
+_rl_cfg.experiment_name = "go2_obstacle_avoidance"
+register_mjlab_task(
+    task_id=OBSTACLE_TASK,
+    env_cfg=unitree_go2_obstacle_env_cfg(),
+    play_env_cfg=unitree_go2_obstacle_env_cfg(play=True),
+    rl_cfg=_rl_cfg,
+    runner_cls=VelocityOnPolicyRunner,
+)

training_server/rl/src/tasks/obstacle_avoidance/env_cfg.py

@@ -0,0 +1,89 @@
+"""81/97-to-12 forward-ray navigation, retaining the Go2 50 Hz position action contract."""
+
+from mjlab.managers import ObservationTermCfg, RewardTermCfg, TerminationTermCfg
+from mjlab.managers.event_manager import EventTermCfg
+from mjlab.sensor import ObjRef, RayCastSensorCfg
+from task_config import (
+    OBSTACLE_TASK,
+    build_terrain_layout,
+    navigation_candidates,
+    validate_task_config,
+)
+
+from src.tasks.velocity.config.go2.env_cfgs import unitree_go2_flat_env_cfg
+
+from . import mdp
+from .terrain import apply_terrain_configuration
+
+
+def apply_obstacle_configuration(cfg, custom, randomize_navigation=True):
+    apply_terrain_configuration(cfg, custom)
+    sensor = custom["sensorCfg"]
+    cfg.scene.sensors = tuple(s for s in cfg.scene.sensors if s.name != "forward_scan") + (
+        RayCastSensorCfg(
+            name="forward_scan",
+            frame=ObjRef(type="body", name="base_link", entity="robot"),
+            pattern=mdp.ForwardFanPatternCfg(fov=sensor["fov"], sensor_mode=sensor["sensorMode"]),
+            ray_alignment="base",
+            max_distance=sensor["maxDistance"],
+            include_geom_groups=(0,),
+            exclude_parent_body=True,
+        ),
+    )
+    for group in ("actor", "critic"):
+        cfg.observations[group].terms["forward_depth"] = ObservationTermCfg(
+            func=mdp.forward_depth, params={"max_distance": sensor["maxDistance"]}
+        )
+        cfg.observations[group].terms["target_error"] = ObservationTermCfg(func=mdp.target_error)
+    layout = build_terrain_layout(custom)
+    size = layout["size"]
+    candidates = navigation_candidates(layout) if randomize_navigation else None
+    cfg.commands["twist"] = mdp.NavigationCommandCfg(
+        resampling_time_range=(1e9, 1e9),
+        goal_offset=tuple(layout["target"][i] - layout["spawn"][i] for i in range(2)),
+        distance_scale=size,
+        **(
+            {
+                "navigation_points": tuple(tuple(point) for point in candidates["points"]),
+                "component_starts": tuple(candidates["componentStarts"]),
+                "component_counts": tuple(candidates["componentCounts"]),
+                "fallback_pairs": tuple(
+                    (tuple(pair[0]), tuple(pair[1])) for pair in candidates["fallbackPairs"]
+                ),
+                "min_goal_distance": candidates["minDistance"],
+            }
+            if candidates
+            else {}
+        ),
+    )
+    if randomize_navigation:
+        cfg.events["reset_base"] = EventTermCfg(func=mdp.reset_navigation_episode, mode="reset")
+    cfg.rewards["goal_progress"] = RewardTermCfg(func=mdp.goal_progress, weight=2.0)
+    cfg.rewards["obstacle_proximity"] = RewardTermCfg(
+        func=mdp.obstacle_proximity,
+        weight=-sensor["avoidanceWeight"],
+        params={
+            "max_distance": sensor["maxDistance"],
+            "safety_distance": sensor["safetyDistance"],
+            **({"floor_size": size} if sensor["sensorMode"] == "multi_ring_raycast" else {}),
+        },
+    )
+    cfg.terminations["outside_map"] = TerminationTermCfg(
+        func=mdp.outside_map, params={"size": size}
+    )
+
+
+def unitree_go2_obstacle_env_cfg(play=False):
+    cfg = unitree_go2_flat_env_cfg(play=False)
+    cfg.sim.nconmax = 128
+    cfg.sim.njmax = 1500
+    cfg.sim.contact_sensor_maxmatch = 500
+    cfg.curriculum = {}
+    for event in ("push_robot", "encoder_bias", "base_com"):
+        cfg.events.pop(event, None)
+    cfg.observations["actor"].enable_corruption = False
+    custom = validate_task_config(OBSTACLE_TASK, {}, 42)
+    apply_obstacle_configuration(cfg, custom)
+    if play:
+        cfg.episode_length_s = 20.0
+    return cfg

training_server/rl/src/tasks/obstacle_avoidance/mdp.py

@@ -0,0 +1,231 @@
+"""Go2 navigation observations: body-frame +X forward, +Y left, +Z up."""
+
+from dataclasses import dataclass
+
+import torch
+from mjlab.managers.command_manager import CommandTerm, CommandTermCfg
+from mjlab.sensor import GridPatternCfg
+
+
+@dataclass
+class ForwardFanPatternCfg(GridPatternCfg):
+    """Inclusive yaw, layer-major pitch; base-aligned with one shared offset."""
+
+    fov: float = 90.0
+    sensor_mode: str = "single_ring_raycast"
+
+    def generate_rays(self, mj_model, device):
+        from task_config import sensor_pattern
+
+        pattern = sensor_pattern(self.sensor_mode, self.fov)
+        yaw = torch.tensor(pattern["yawAngles"], device=device) * torch.pi / 180
+        pitch = torch.tensor(pattern["pitchAngles"], device=device) * torch.pi / 180
+        p, y = torch.meshgrid(pitch, yaw, indexing="ij")
+        directions = torch.stack((p.cos() * y.cos(), p.cos() * y.sin(), p.sin()), dim=-1).reshape(
+            -1, 3
+        )
+        offsets = torch.tensor([0.3, 0, 0.05], device=device).repeat(pattern["rayCount"], 1)
+        return offsets, directions
+
+
+@dataclass(kw_only=True)
+class NavigationCommandCfg(CommandTermCfg):
+    goal_offset: tuple[float, float] = (10.0, 0.0)
+    distance_scale: float = 12.0
+    speed: float = 0.6
+    arrival_radius: float = 0.5
+    navigation_points: tuple[tuple[float, float], ...] = ()
+    component_starts: tuple[int, ...] = ()
+    component_counts: tuple[int, ...] = ()
+    fallback_pairs: tuple[tuple[tuple[float, float], tuple[float, float]], ...] = ()
+    min_goal_distance: float = 2.0
+
+    def build(self, env):
+        return NavigationCommand(self, env)
+
+
+class NavigationCommand(CommandTerm):
+    cfg: NavigationCommandCfg
+
+    def __init__(self, cfg, env):
+        super().__init__(cfg, env)
+        self.metrics["distance_to_goal"] = torch.zeros(self.num_envs, device=self.device)
+        self.goals_w = self._env.scene.env_origins[:, :2] + torch.tensor(
+            self.cfg.goal_offset, device=self.device
+        )
+        self._points = (
+            torch.tensor(self.cfg.navigation_points, dtype=torch.float32, device=self.device)
+            if self.cfg.navigation_points
+            else None
+        )
+        self._component_starts = torch.tensor(
+            self.cfg.component_starts, dtype=torch.long, device=self.device
+        )
+        self._component_counts = torch.tensor(
+            self.cfg.component_counts, dtype=torch.long, device=self.device
+        )
+        self._fallback_pairs = (
+            torch.tensor(self.cfg.fallback_pairs, dtype=torch.float32, device=self.device)
+            if self.cfg.fallback_pairs
+            else None
+        )
+
+    def errors(self):
+        robot = self._env.scene["robot"]
+        delta = self.goals_w - robot.data.root_link_pos_w[:, :2]
+        q = robot.data.root_link_quat_w
+        yaw = torch.atan2(
+            2 * (q[:, 0] * q[:, 3] + q[:, 1] * q[:, 2]),
+            1 - 2 * (q[:, 2].square() + q[:, 3].square()),
+        )
+        heading = torch.atan2(delta[:, 1], delta[:, 0]) - yaw
+        heading = torch.atan2(heading.sin(), heading.cos())
+        distance = delta.norm(dim=1)
+        heading = torch.where(distance < self.cfg.arrival_radius, 0.0, heading)
+        return heading, distance, delta
+
+    @property
+    def command(self):
+        # Compute from current physics, not a one-step-old command manager cache.
+        heading, distance, _ = self.errors()
+        moving = distance >= self.cfg.arrival_radius
+        vx = self.cfg.speed * heading.cos().clamp(min=0)
+        return torch.stack(
+            (
+                torch.where(moving, vx, 0.0),
+                torch.zeros_like(vx),
+                torch.where(moving, heading.clamp(-1, 1), 0.0),
+            ),
+            dim=1,
+        )
+
+    def _update_metrics(self):
+        self.metrics["distance_to_goal"][:] = self.errors()[1]
+
+    def sample_episode(self, env_ids):
+        """Randomize a reachable free-space pair and the initial heading."""
+        if self._points is None or len(env_ids) == 0:
+            return
+        count = len(env_ids)
+        component = torch.randint(len(self._component_starts), (count,), device=self.device)
+        starts = self._component_starts[component]
+        counts = self._component_counts[component]
+
+        def indices():
+            return starts + (torch.rand(count, device=self.device) * counts).long()
+
+        spawn = self._points[indices()]
+        target = self._points[indices()]
+        valid = (target - spawn).norm(dim=1) >= self.cfg.min_goal_distance
+        # Keep this branch-free on CUDA; unresolved samples use a guaranteed component pair.
+        for _ in range(16):
+            candidate = self._points[indices()]
+            target = torch.where(valid[:, None], target, candidate)
+            valid = (target - spawn).norm(dim=1) >= self.cfg.min_goal_distance
+        fallback = self._fallback_pairs[component]
+        spawn = torch.where(valid[:, None], spawn, fallback[:, 0])
+        target = torch.where(valid[:, None], target, fallback[:, 1])
+
+        robot = self._env.scene["robot"]
+        default = robot.data.default_root_state[env_ids].clone()
+        position = torch.cat((spawn, default[:, 2:3]), dim=1)
+        yaw = torch.rand(count, device=self.device) * (2 * torch.pi) - torch.pi
+        half = yaw / 2
+        quaternion = torch.stack(
+            (half.cos(), torch.zeros_like(half), torch.zeros_like(half), half.sin()), dim=1
+        )
+        robot.write_root_link_pose_to_sim(torch.cat((position, quaternion), dim=1), env_ids=env_ids)
+        robot.write_root_link_velocity_to_sim(torch.zeros_like(default[:, 7:13]), env_ids=env_ids)
+        self.goals_w[env_ids] = target
+
+    def _resample_command(self, env_ids):
+        # Episode pairs are sampled by the reset event before manager resets.
+        self.metrics["distance_to_goal"][env_ids] = self.errors()[1][env_ids]
+
+    def _update_command(self):
+        # No cached command: the property is evaluated after every physics update.
+        self._update_metrics()
+
+
+def reset_navigation_episode(env, env_ids):
+    if env_ids is None:
+        env_ids = torch.arange(env.num_envs, device=env.device, dtype=torch.int)
+    env.command_manager.get_term("twist").sample_episode(env_ids)
+
+
+def forward_depth(env, sensor_name="forward_scan", max_distance=4.0):
+    distances = env.scene[sensor_name].data.distances
+    distances = torch.where(distances < 0, max_distance, distances)
+    return distances.clamp(0, max_distance) / max_distance
+
+
+def target_error(env):
+    command = env.command_manager.get_term("twist")
+    heading, distance, _ = command.errors()
+    return torch.stack(
+        (heading / torch.pi, (distance / command.cfg.distance_scale).clamp(0, 1)), dim=1
+    )
+
+
+def standard_floor_id(model, size):
+    """Fail closed: match the compiled first layout box, not a name alone."""
+    import mujoco
+    import numpy as np
+
+    data = mujoco.MjData(model)
+    mujoco.mj_forward(model, data)
+    ids = [
+        i
+        for i in range(model.ngeom)
+        if model.geom_type[i] == 6
+        and model.body_weldid[model.geom_bodyid[i]] == 0
+        and model.geom_group[i] == 0
+        and np.allclose(data.geom_xpos[i], [0, 0, -0.1], atol=1e-7, rtol=0)
+        and np.allclose(model.geom_size[i], [size / 2, size / 2, 0.1], atol=1e-7, rtol=0)
+        and np.allclose(data.geom_xmat[i], [1, 0, 0, 0, 1, 0, 0, 0, 1], atol=1e-7, rtol=0)
+    ]
+    if len(ids) != 1 or model.geom(ids[0]).name != "terrain_0":
+        raise ValueError("multi-ring reward requires one verified static standard floor")
+    return ids[0]
+
+
+def floor_top_hits(data, size):
+    # Warp ray hits use the same world-centered coordinates in each independent env.
+    # Float32 tolerance: 10 micrometres; never erase 0.05m low obstacles.
+    hit = data.hit_pos_w
+    normal = data.normals_w
+    return (
+        (data.distances >= 0)
+        & torch.isfinite(data.distances)
+        & (hit[..., :2].abs() <= size / 2 + 1e-5).all(dim=-1)
+        & (hit[..., 2].abs() <= 1e-5)
+        & (normal[..., :2].abs() <= 1e-5).all(dim=-1)
+        & ((normal[..., 2] - 1).abs() <= 1e-5)
+    )
+
+
+def obstacle_proximity(env, max_distance=4.0, safety_distance=0.5, floor_size=None):
+    depths = forward_depth(env, max_distance=max_distance)
+    if floor_size is not None:
+        # Validate once per environment/model. Observation data remains untouched.
+        if not hasattr(env, "_multi_ring_floor_id"):
+            env._multi_ring_floor_id = standard_floor_id(env.sim.mj_model, floor_size)
+        depths = torch.where(
+            floor_top_hits(env.scene["forward_scan"].data, floor_size), 1.0, depths
+        )
+    closest = depths.amin(dim=1) * max_distance
+    return ((safety_distance - closest).clamp(min=0) / safety_distance).square()
+
+
+def goal_progress(env):
+    command = env.command_manager.get_term("twist")
+    _, distance, delta = command.errors()
+    velocity = env.scene["robot"].data.root_link_lin_vel_w[:, :2]
+    progress = (velocity * delta / distance.clamp(min=1e-6).unsqueeze(1)).sum(dim=1)
+    return torch.where(distance >= command.cfg.arrival_radius, progress, 0.0)
+
+
+def outside_map(env, size=12.0):
+    # Every Warp environment is an independent copy of the same world-centered map.
+    position = env.scene["robot"].data.root_link_pos_w[:, :2]
+    return (position.abs() > size / 2 - 0.3).any(dim=1)

training_server/rl/src/tasks/obstacle_avoidance/terrain.py

@@ -0,0 +1,70 @@
+"""Install the exact bounded boxes-v1 deployment layout in mjlab's generator."""
+
+from dataclasses import dataclass, field
+
+import mujoco
+import numpy as np
+from mjlab.envs import mdp
+from mjlab.managers.event_manager import EventTermCfg
+from mjlab.terrains import TerrainEntityCfg, TerrainGeneratorCfg
+from mjlab.terrains.terrain_generator import SubTerrainCfg, TerrainGeometry, TerrainOutput
+from task_config import build_terrain_layout
+
+
+@dataclass
+class DeploymentTerrainCfg(SubTerrainCfg):
+    layout: dict = field(default_factory=dict)
+
+    def function(self, difficulty, spec, rng):
+        # Generator translates patch-local corner coordinates by (-size/2,-size/2).
+        half = self.layout["size"] / 2
+        translation = np.array([half, half, 0])
+        body = spec.body("terrain")
+        geometries = []
+        for box in self.layout["boxes"]:
+            geom = body.add_geom(
+                type=mujoco.mjtGeom.mjGEOM_BOX,
+                pos=np.array(box["pos"]) + translation,
+                size=box["size"],
+                group=0,
+                contype=1,
+                conaffinity=1,
+                friction=[self.layout["friction"], 0.005, 0.0001],
+                priority=1,
+                condim=3,
+            )
+            geometries.append(TerrainGeometry(geom=geom))
+        spawn = np.array(self.layout["spawn"])
+        spawn[2] = 0  # Robot init_state contributes the 0.32 m body height.
+        return TerrainOutput(origin=spawn + translation, geometries=geometries)
+
+
+def apply_terrain_configuration(cfg, custom):
+    layout = build_terrain_layout(custom)
+    size = layout["size"]
+    cfg.scene.entities["robot"].init_state.pos = (0.0, 0.0, layout["spawn"][2])
+    cfg.scene.entities["robot"].init_state.rot = tuple(layout["spawnQuaternion"])
+    cfg.scene.terrain = TerrainEntityCfg(
+        terrain_type="generator",
+        max_init_terrain_level=0,
+        terrain_generator=TerrainGeneratorCfg(
+            seed=custom["seed"],
+            size=(size, size),
+            num_rows=1,
+            num_cols=1,
+            border_width=0,
+            curriculum=False,
+            color_scheme="none",
+            sub_terrains={"deployment": DeploymentTerrainCfg(layout=layout)},
+        ),
+    )
+    cfg.curriculum.pop("terrain_levels", None)
+    cfg.events.pop("randomize_terrain", None)
+    # Install the deterministic reset first; the navigation task may replace it with
+    # collision-free random point-goal sampling after the command term is configured.
+    cfg.events["reset_base"] = EventTermCfg(
+        func=mdp.reset_root_state_uniform,
+        mode="reset",
+        params={"pose_range": {}, "velocity_range": {}},
+    )
+    cfg.events["foot_friction"].params["ranges"] = (layout["friction"], layout["friction"])

training_server/rl/src/tasks/velocity/rl/runner.py

@@ -1,4 +1,6 @@
+import json
 import os
+from pathlib import Path
 
 import wandb
 
@@ -23,6 +25,19 @@ class VelocityOnPolicyRunner(MjlabOnPolicyRunner):
     )  # type: ignore[assignment]
     onnx_path = os.path.join(policy_path, filename)
     metadata = get_base_metadata(self.env.unwrapped, run_name)
+    deployment = getattr(self.env.unwrapped, "platform_deployment", None)
+    if deployment is not None:
+      metadata["platform_deployment"] = json.dumps(deployment, separators=(",", ":"))
+      Path(policy_path, "deployment.json").write_text(
+        json.dumps(deployment, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
+      )
+    initialization = getattr(self.env.unwrapped, "platform_initialization", None)
+    if initialization is not None:
+      # Only content identity and protocol, never local source paths or arbitrary infos.
+      from pretrained import public_initialization_metadata
+      metadata["pretrained_initialization"] = json.dumps(
+        public_initialization_metadata(initialization), separators=(",", ":")
+      )
     attach_metadata_to_onnx(onnx_path, metadata)
     if self.logger.logger_type in ["wandb"]:
       wandb.save(policy_path + filename, base_path=os.path.dirname(policy_path))