feat(training): release V0.8 自调参 Agent

2026-09-02 13:49:34 +08:00
parent cffac29a03
commit deead17a9a
47 changed files with 4986 additions and 96 deletions
@@ -0,0 +1,222 @@
+"""Deterministic, headless evaluation for Unitree Go2 velocity policies."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import sys
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from statistics import fmean, pstdev
+from typing import Literal
+
+TRAINER_ROOT = Path(__file__).resolve().parents[1]
+SERVICE_ROOT = TRAINER_ROOT.parent
+for source_root in (TRAINER_ROOT, SERVICE_ROOT):
+  if str(source_root) not in sys.path:
+    sys.path.insert(0, str(source_root))
+
+import torch
+import tyro
+import warp as wp
+
+if not hasattr(wp, "context"):
+  from warp._src import context as warp_context
+
+  wp.context = warp_context  # type: ignore[attr-defined]
+
+from mjlab.envs import ManagerBasedRlEnv
+from mjlab.rl import MjlabOnPolicyRunner, RslRlVecEnvWrapper
+from mjlab.tasks.registry import list_tasks, load_env_cfg, load_rl_cfg, load_runner_cls
+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 = (
+  (0.0, 0.0, 0.0),
+  (0.5, 0.0, 0.0),
+  (1.0, 0.0, 0.0),
+  (1.5, 0.0, 0.0),
+  (0.0, 0.5, 0.0),
+  (0.0, -0.5, 0.0),
+  (0.0, 0.0, 0.5),
+  (0.0, 0.0, -0.5),
+  (0.8, 0.25, 0.35),
+)
+METRIC_NAMES = (
+  "linear_velocity_rmse",
+  "angular_velocity_rmse",
+  "mean_action_acc",
+  "orientation_error",
+  "slip_velocity",
+  "mechanical_power",
+)
+
+
+@dataclass(frozen=True)
+class EvaluateConfig:
+  checkpoint: str
+  output: str
+  reward_config: str | None = None
+  num_envs: int = 256
+  steps_per_seed: int = 1000
+  seeds: tuple[int, ...] = field(default_factory=lambda: (101, 202, 303))
+  device: str | None = None
+  gpu_ids: list[int] | Literal["all"] | None = field(default_factory=lambda: [0])
+
+
+def _sha256(path: Path) -> str:
+  digest = hashlib.sha256()
+  with path.open("rb") as stream:
+    for chunk in iter(lambda: stream.read(1024 * 1024), b""):
+      digest.update(chunk)
+  return digest.hexdigest()
+
+
+def _fixed_commands(command: torch.Tensor) -> torch.Tensor:
+  values = torch.as_tensor(SCENARIOS, device=command.device, dtype=command.dtype)
+  indexes = torch.arange(command.shape[0], device=command.device) % values.shape[0]
+  command[:] = values[indexes]
+  return command
+
+
+def _evaluate_seed(task_id: str, cfg: EvaluateConfig, seed: int) -> dict[str, float]:
+  torch.manual_seed(seed)
+  env_cfg = load_env_cfg(task_id, play=False)
+  agent_cfg = load_rl_cfg(task_id)
+  env_cfg.seed = seed
+  env_cfg.scene.num_envs = cfg.num_envs
+  env_cfg.curriculum = {}
+  env_cfg.observations["actor"].enable_corruption = False
+  env_cfg.events.pop("push_robot", None)
+  twist_cfg = env_cfg.commands["twist"]
+  assert isinstance(twist_cfg, UniformVelocityCommandCfg)
+  twist_cfg.heading_command = False
+  twist_cfg.ranges.heading = None
+  twist_cfg.rel_heading_envs = 0.0
+  twist_cfg.rel_standing_envs = 0.0
+  twist_cfg.resampling_time_range = (1.0e9, 1.0e9)
+  if cfg.reward_config:
+    reward_path = Path(cfg.reward_config).expanduser().resolve(strict=True)
+    with reward_path.open(encoding="utf-8") as stream:
+      apply_reward_configuration(env_cfg, validate_configuration(json.load(stream)))
+
+  env = ManagerBasedRlEnv(cfg=env_cfg, device=cfg.device or "cuda:0")
+  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)
+    runner.load(
+      str(Path(cfg.checkpoint).expanduser().resolve(strict=True)),
+      load_cfg={"actor": True},
+      strict=True,
+      map_location=str(wrapped.device),
+    )
+    policy = runner.get_inference_policy(device=str(wrapped.device))
+    twist = wrapped.unwrapped.command_manager.get_term("twist")
+    _fixed_commands(twist.command)
+    obs = wrapped.get_observations()
+
+    sums = {name: 0.0 for name in METRIC_NAMES}
+    samples = 0
+    terminations = 0
+    completions = 0
+    with torch.inference_mode():
+      for _ in range(cfg.steps_per_seed):
+        _fixed_commands(twist.command)
+        obs = wrapped.get_observations()
+        actions = policy(obs)
+        obs, _rewards, _dones, _extras = wrapped.step(actions)
+        manager = wrapped.unwrapped.metrics_manager
+        for index, name in enumerate(manager.active_terms):
+          if name in sums:
+            sums[name] += float(torch.sum(manager._step_values[:, index]).item())
+        samples += wrapped.num_envs
+        terminated = wrapped.unwrapped.reset_terminated
+        timed_out = wrapped.unwrapped.reset_time_outs
+        terminations += int(torch.count_nonzero(terminated).item())
+        completions += int(torch.count_nonzero(terminated | timed_out).item())
+    result = {name: sums[name] / max(samples, 1) for name in METRIC_NAMES}
+    result["fall_rate"] = terminations / max(completions, wrapped.num_envs)
+    return result
+  finally:
+    wrapped.close()
+
+
+def run_evaluation(task_id: str, cfg: EvaluateConfig) -> dict:
+  configure_torch_backends()
+  selected = cfg.gpu_ids
+  if selected is None:
+    os.environ["CUDA_VISIBLE_DEVICES"] = ""
+    device = "cpu"
+  else:
+    if selected == "all":
+      selected = [0]
+    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, selected))
+    device = cfg.device or "cuda:0"
+  os.environ["MUJOCO_GL"] = "egl"
+  cfg = EvaluateConfig(**{**asdict(cfg), "device": device})
+
+  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)
+    ],
+  }
+  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:
+    for name, value in metrics.items():
+      writer.add_scalar(f"Evaluation/{name}", value, 0)
+  finally:
+    writer.close()
+  print("MUJOCO_EVALUATION " + json.dumps({"output": str(output), "metrics": metrics}))
+  return result
+
+
+def main() -> None:
+  import mjlab.tasks  # noqa: F401
+  import src.tasks  # noqa: F401
+
+  chosen_task, remaining = tyro.cli(
+    tyro.extras.literal_type_from_choices(list_tasks()),
+    add_help=False,
+    return_unknown_args=True,
+    config=mjlab.TYRO_FLAGS,
+  )
+  args = tyro.cli(
+    EvaluateConfig,
+    args=remaining,
+    prog=sys.argv[0] + f" {chosen_task}",
+    config=mjlab.TYRO_FLAGS,
+  )
+  run_evaluation(chosen_task, args)
+
+
+if __name__ == "__main__":
+  main()