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

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(