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

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)