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

training_server/rl/scripts/evaluate.py

@@ -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()

training_server/rl/scripts/train.py

@@ -1,5 +1,6 @@
 """Script to train RL agent with RSL-RL."""
 
+import json
 import logging
 import os
 import sys
@@ -10,8 +11,10 @@ from typing import Literal, cast
 
 # 训练器作为仓库内置子集直接从 scripts/ 启动，不要求额外执行 pip install -e。
 TRAINER_ROOT = Path(__file__).resolve().parents[1]
-if str(TRAINER_ROOT) not in sys.path:
-  sys.path.insert(0, str(TRAINER_ROOT))
+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 tyro
 import warp as wp
@@ -32,6 +35,8 @@ 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 tuning.schema import apply_reward_configuration, validate_configuration
+
 
 @dataclass(frozen=True)
 class TrainConfig:
@@ -44,6 +49,10 @@ class TrainConfig:
   enable_nan_guard: bool = False
   torchrunx_log_dir: str | None = None
   gpu_ids: list[int] | Literal["all"] | None = field(default_factory=lambda: [0])
+  output_dir: str | None = None
+  resume_checkpoint: str | None = None
+  reward_config: str | None = None
+  reward_config_json: str | None = None
 
   @staticmethod
   def from_task(task_id: str) -> "TrainConfig":
@@ -52,7 +61,27 @@ class TrainConfig:
     return TrainConfig(env=env_cfg, agent=agent_cfg)
 
 
+def _load_reward_config(path: str | None, inline: str | None) -> 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))
+  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))
+
+
 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)
+
   cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
   if cuda_visible == "":
     device = "cpu"
@@ -109,11 +138,14 @@ def run_train(task_id: str, cfg: TrainConfig, log_dir: Path) -> None:
   log_root_path = log_dir.parent  # Go up from specific run dir to experiment dir.
 
   resume_path: Path | None = None
-  if cfg.agent.resume:
-      # Load checkpoint from local filesystem.
-      resume_path = get_checkpoint_path(
-        log_root_path, cfg.agent.load_run, cfg.agent.load_checkpoint
-      )
+  explicit_resume = cfg.resume_checkpoint is not None
+  if explicit_resume:
+    resume_path = Path(cfg.resume_checkpoint).expanduser().resolve(strict=True)
+  elif cfg.agent.resume:
+    # Load checkpoint from local filesystem.
+    resume_path = get_checkpoint_path(
+      log_root_path, cfg.agent.load_run, cfg.agent.load_checkpoint
+    )
 
   # Only record videos on rank 0 to avoid multiple workers writing to the same files.
   if cfg.video and rank == 0:
@@ -141,16 +173,32 @@ def run_train(task_id: str, cfg: TrainConfig, log_dir: Path) -> None:
   runner.add_git_repo_to_log(__file__)
   if resume_path is not None:
     print(f"[INFO]: Loading model checkpoint from: {resume_path}")
-    runner.load(str(resume_path))
+    runner.load(str(resume_path), map_location=device)
+    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.
+      runner.current_learning_iteration += 1
 
   # Only write config files from rank 0 to avoid race conditions.
   if rank == 0:
     dump_yaml(log_dir / "params" / "env.yaml", env_cfg)
     dump_yaml(log_dir / "params" / "agent.yaml", agent_cfg)
+    if reward_config is not None:
+      reward_snapshot = log_dir / "params" / "reward_config.json"
+      reward_snapshot.parent.mkdir(parents=True, exist_ok=True)
+      reward_snapshot.write_text(
+        json.dumps(reward_config, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
+      )
 
-  runner.learn(
-    num_learning_iterations=cfg.agent.max_iterations, init_at_random_ep_len=True
+  iterations = cfg.agent.max_iterations
+  if explicit_resume:
+    iterations = max(0, cfg.agent.max_iterations - runner.current_learning_iteration)
+  print(
+    f"[INFO] Learning target: current={runner.current_learning_iteration}, "
+    f"additional={iterations}, target={cfg.agent.max_iterations}",
+    flush=True,
   )
+  runner.learn(num_learning_iterations=iterations, init_at_random_ep_len=True)
 
   env.close()
 
@@ -159,12 +207,15 @@ def launch_training(task_id: str, args: TrainConfig | None = None):
   args = args or TrainConfig.from_task(task_id)
 
   # Create log directory once before launching workers.
-  log_root_path = Path("logs") / "rsl_rl" / args.agent.experiment_name
-  log_root_path.resolve()
-  log_dir_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
-  if args.agent.run_name:
-    log_dir_name += f"_{args.agent.run_name}"
-  log_dir = log_root_path / log_dir_name
+  if args.output_dir:
+    log_dir = Path(args.output_dir).expanduser().resolve()
+  else:
+    log_root_path = (Path("logs") / "rsl_rl" / args.agent.experiment_name).resolve()
+    log_dir_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    if args.agent.run_name:
+      log_dir_name += f"_{args.agent.run_name}"
+    log_dir = log_root_path / log_dir_name
+  log_dir.mkdir(parents=True, exist_ok=True)
 
   # Select GPUs based on CUDA_VISIBLE_DEVICES and user specification.
   selected_gpus, num_gpus = select_gpus(args.gpu_ids)