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

training_server/server.py

@@ -25,14 +25,23 @@ from pathlib import Path
 from typing import Any
 from urllib.parse import parse_qs, unquote, urlsplit
 
+from pretrained_sources import PretrainedSources, SourceError
+from task_config import (
+    OBSTACLE_TASK,
+    TaskConfigError,
+    deployment_metadata,
+    task_metadata,
+    validate_task_config,
+)
 from tuning.manager import TuningError, TuningManager
 from tuning.process import GpuLease, ResourceBusyError
-from tuning.schema import RewardConfigError
+from tuning.schema import RewardConfigError, validate_configuration
 from tuning.scoring import EvaluationError
 
 VERSION = "0.4.0"
-# 浏览器当前 ONNX 运行时只实现 Go2 的 47→12 部署契约；其他任务须由服务启动参数显式放行。
-DEFAULT_TASKS = ("Unitree-Go2-Flat",)
+MAX_REQUEST_BYTES = 128 * 1024  # Bounded full boxes-v1 payload (<=257 boxes).
+# Rough 可以训练，但其高度扫描 actor 不允许冒充浏览器 Flat 部署。
+DEFAULT_TASKS = ("Unitree-Go2-Flat", "Unitree-Go2-Rough", OBSTACLE_TASK)
 ACTIVE_STATES = {"queued", "running"}
 MAX_JOBS = 20
 ANSI_ESCAPE = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]")
@@ -71,6 +80,12 @@ class TrainingConfig:
     gpu_ids: list[int]
     wandb_mode: str
     reward_config: dict[str, Any] | None = None
+    terrain_preset: str | None = None
+    terrain_params: dict[str, Any] = field(default_factory=dict)
+    sensor_cfg: dict[str, Any] | None = None
+    task_config: dict[str, Any] | None = None
+    deployment: dict[str, Any] = field(default_factory=dict)
+    pretrained: dict[str, Any] | None = None
 
 
 @dataclass
@@ -96,6 +111,8 @@ class TrainingJob:
             "id": self.id,
             "state": self.state,
             "taskId": self.config.task_id,
+            "deployment": self.config.deployment,
+            "pretrained": self.config.pretrained,
             "createdAt": self.created_at,
             "startedAt": self.started_at,
             "endedAt": self.ended_at,
@@ -117,6 +134,7 @@ class TrainingManager:
         tasks: tuple[str, ...],
         check_environment: bool = True,
         lease: GpuLease | None = None,
+        sources: PretrainedSources | None = None,
     ):
         self.trainer_root = trainer_root.expanduser().resolve()
         self.python = str(Path(python).expanduser()) if os.sep in python else python
@@ -127,6 +145,7 @@ class TrainingManager:
         self._environment_error: str | None | bool = False
         self.lease = lease or GpuLease()
         self.preset_resolver: Any = None
+        self.sources = sources
 
     def readiness_error(self) -> str | None:
         if not self.trainer_root.is_dir():
@@ -172,6 +191,13 @@ class TrainingManager:
             "trainerRoot": str(self.trainer_root),
             "python": self.python,
             "tasks": list(self.tasks),
+            "pretrainedSources": self.sources.catalog() if self.sources else [],
+            "pretrainedUpload": {
+                "enabled": self.sources is not None, "templateId": "go2-legacy47-v1",
+                "formats": {"pt": 256 * 1024**2, "onnx": 64 * 1024**2},
+                "endpoint": "/api/training/pretrained-sources/upload",
+            },
+            "taskMetadata": task_metadata(self.tasks),
             "activeJobId": self.active_job_id(),
             "error": error,
         }
@@ -179,6 +205,25 @@ class TrainingManager:
     def parse_config(self, payload: Any) -> TrainingConfig:
         if not isinstance(payload, dict):
             raise ApiError(HTTPStatus.BAD_REQUEST, "请求体必须是 JSON 对象")
+        allowed = {
+            "taskId",
+            "numEnvs",
+            "maxIterations",
+            "seed",
+            "runName",
+            "device",
+            "gpuIds",
+            "wandbMode",
+            "rewardPresetId",
+            "pretrainedSourceId",
+            "terrainPreset",
+            "terrainParams",
+            "sensorCfg",
+            "sensorType",
+            "customTerrainBoxes",
+        }
+        if payload.keys() - allowed:
+            raise ApiError(HTTPStatus.BAD_REQUEST, "请求包含未知字段（不接受配置路径/MJCF）")
         task_id = payload.get("taskId")
         if task_id not in self.tasks:
             raise ApiError(HTTPStatus.BAD_REQUEST, f"不允许的训练任务：{task_id}")
@@ -216,19 +261,46 @@ class TrainingManager:
         preset_id = payload.get("rewardPresetId")
         reward_config = None
         if preset_id is not None:
+            if task_id != "Unitree-Go2-Flat":
+                raise ApiError(HTTPStatus.BAD_REQUEST, "自动调参奖励 preset 仅支持 Flat 任务")
             if not isinstance(preset_id, str) or not re.fullmatch(r"[0-9a-f]{32}", preset_id):
                 raise ApiError(HTTPStatus.BAD_REQUEST, "rewardPresetId 格式无效")
             if self.preset_resolver is None:
                 raise ApiError(HTTPStatus.BAD_REQUEST, "奖励 preset 服务未就绪")
             try:
-                reward_config = self.preset_resolver(preset_id)
+                reward_config = validate_configuration(
+                    self.preset_resolver(preset_id, task_id), task_id
+                )
             except KeyError as error:
                 raise ApiError(HTTPStatus.BAD_REQUEST, "奖励 preset 不存在") from error
+            except RewardConfigError as error:
+                raise ApiError(HTTPStatus.BAD_REQUEST, str(error)) from error
+        seed = integer("seed", 0, 2_147_483_647)
+        try:
+            custom = validate_task_config(task_id, payload, seed)
+        except TaskConfigError as error:
+            raise ApiError(HTTPStatus.BAD_REQUEST, str(error)) from error
+        pretrained = None
+        if "pretrainedSourceId" in payload:
+            if self.sources is None:
+                raise ApiError(
+                    HTTPStatus.BAD_REQUEST, "服务尚未注册基础策略，请配置--pretrained-sources"
+                )
+            try:
+                pretrained = self.sources.bind(payload["pretrainedSourceId"], task_id, custom)
+            except SourceError as error:
+                raise ApiError(HTTPStatus.BAD_REQUEST, str(error)) from error
         return TrainingConfig(
+            pretrained=pretrained,
             task_id=task_id,
+            terrain_preset=custom["terrainPreset"] if custom else None,
+            terrain_params=custom["terrainParams"] if custom else {},
+            sensor_cfg=custom["sensorCfg"] if custom else None,
+            task_config=custom,
+            deployment=deployment_metadata(task_id, custom, seed),
             num_envs=integer("numEnvs", 1, 16384),
             max_iterations=integer("maxIterations", 1, 1_000_000),
-            seed=integer("seed", 0, 2_147_483_647),
+            seed=seed,
             run_name=run_name,
             device=device,
             gpu_ids=raw_gpu_ids,
@@ -325,7 +397,9 @@ class TrainingManager:
                     os.killpg(process.pid, signal.SIGKILL)
                 process.wait(timeout=2)
 
-    def command_for(self, config: TrainingConfig) -> list[str]:
+    def command_for(
+        self, config: TrainingConfig, task_config_path: Path | None = None
+    ) -> list[str]:
         command = [
             self.python,
             "-u",
@@ -336,6 +410,12 @@ class TrainingManager:
             f"--agent.seed={config.seed}",
             f"--agent.run-name={config.run_name}",
         ]
+        if task_config_path is not None:
+            command.extend(("--task-config", str(task_config_path)))
+        if config.pretrained is not None:
+            if self.sources is None:
+                raise SourceError("基础策略快照服务未配置，拒绝随机初始化")
+            command.extend(self.sources.arguments(config.pretrained))
         if config.reward_config is not None:
             command.extend(
                 (
@@ -387,13 +467,21 @@ class TrainingManager:
 
     def _run(self, job: TrainingJob) -> None:
         before = self._artifact_snapshot()
-        command = self.command_for(job.config)
         environment = os.environ.copy()
         # 默认离线记录，保留本地 W&B 指标但不要求 API Key；只有前端明确选择
         # online 时才允许 wandb 发起登录/联网。
         environment["WANDB_MODE"] = job.config.wandb_mode
         environment["WANDB_SILENT"] = "true"
         try:
+            config_path = None
+            if job.config.task_config is not None:
+                job_dir = self.trainer_root / "logs" / "rsl_rl" / "web_jobs" / job.id
+                job_dir.mkdir(parents=True, exist_ok=True)
+                config_path = job_dir / "training_config.json"
+                config_path.write_text(json.dumps(job.config.task_config), encoding="utf-8")
+            command = self.command_for(job.config, config_path)
+            if config_path is not None:
+                command.extend(("--output-dir", str(config_path.parent)))
             # Popen 与 process 登记必须和取消检查处于同一个临界区：cancel() 要么在
             # 创建前标记取消，要么在创建后取得进程并终止，不能落入二者之间。
             with self.lock:
@@ -500,7 +588,7 @@ class TrainingRequestHandler(BaseHTTPRequestHandler):
             self._json(HTTPStatus.NOT_FOUND, {"error": "调参 session、trial 或 proposal 不存在"})
         elif isinstance(error, ResourceBusyError):
             self._json(HTTPStatus.CONFLICT, {"error": str(error)})
-        elif isinstance(error, (TuningError, RewardConfigError, EvaluationError)):
+        elif isinstance(error, (TuningError, RewardConfigError, EvaluationError, SourceError)):
             self._json(HTTPStatus.BAD_REQUEST, {"error": str(error)})
         else:
             self._json(
@@ -523,15 +611,47 @@ class TrainingRequestHandler(BaseHTTPRequestHandler):
             length = int(self.headers.get("Content-Length", "0"))
         except ValueError as error:
             raise ApiError(HTTPStatus.BAD_REQUEST, "Content-Length 无效") from error
-        if length <= 0 or length > 32 * 1024:
+        if length <= 0 or length > MAX_REQUEST_BYTES:
             raise ApiError(
-                HTTPStatus.REQUEST_ENTITY_TOO_LARGE, "训练请求体不能为空且不能超过 32 KiB"
+                HTTPStatus.REQUEST_ENTITY_TOO_LARGE, "训练请求体不能为空且不能超过 128 KiB"
             )
         try:
             return json.loads(self.rfile.read(length))
         except (UnicodeDecodeError, json.JSONDecodeError) as error:
             raise ApiError(HTTPStatus.BAD_REQUEST, "训练请求不是有效 JSON") from error
 
+    def _upload(self):
+        # This dedicated binary route must never pass through the 128KiB JSON reader.
+        self.close_connection = True  # no unread extra bytes may become a second request
+        if self.manager.sources is None:
+            raise ApiError(HTTPStatus.SERVICE_UNAVAILABLE, "基础策略上传服务不可用")
+        if self.headers.get("Transfer-Encoding") or self.headers.get("Content-Encoding"):
+            raise ApiError(HTTPStatus.BAD_REQUEST, "上传不接受chunked/压缩编码")
+        lengths = self.headers.get_all("Content-Length", [])
+        if len(lengths) != 1 or not re.fullmatch(r"[0-9]{1,10}", lengths[0]):
+            raise ApiError(HTTPStatus.BAD_REQUEST, "上传必须提供唯一Content-Length")
+        length = int(lengths[0])
+        if self.headers.get("Content-Type") != "application/octet-stream":
+            raise ApiError(HTTPStatus.BAD_REQUEST, "上传Content-Type必须是application/octet-stream")
+        query = parse_qs(urlsplit(self.path).query, keep_blank_values=True)
+        if set(query) != {"format", "template", "name"} or any(len(v) != 1 for v in query.values()):
+            raise ApiError(HTTPStatus.BAD_REQUEST, "上传必须显式提供format/template/name")
+        fmt, template, name = (query[key][0] for key in ("format", "template", "name"))
+        limit = {"pt": 256 * 1024**2, "onnx": 64 * 1024**2}.get(fmt)
+        if limit is not None and (length == 0 or length > limit):
+            raise ApiError(HTTPStatus.REQUEST_ENTITY_TOO_LARGE, "上传为空或超过格式大小上限")
+        if len(name) > 255:
+            raise ApiError(HTTPStatus.BAD_REQUEST, "上传显示名称过长")
+        try:
+            result = self.manager.sources.receive_upload(
+                self.rfile, length, fmt, template, name, set_timeout=self.connection.settimeout,
+            )
+        except OSError as error:
+            raise ApiError(
+                HTTPStatus.SERVICE_UNAVAILABLE, "上传存储写入失败，请检查本机磁盘空间/权限"
+            ) from error
+        self._json(HTTPStatus.CREATED, result)
+
     @staticmethod
     def _route(path: str) -> tuple[str | None, bool]:
         match = re.fullmatch(r"/api/training/jobs/([0-9a-f]{32})(/artifacts/policy\.onnx)?", path)
@@ -628,6 +748,9 @@ class TrainingRequestHandler(BaseHTTPRequestHandler):
         try:
             self._ensure_request()
             path = urlsplit(self.path).path
+            if path == "/api/training/pretrained-sources/upload":
+                self._upload()
+                return
             if path == "/api/training/jobs":
                 self._json(HTTPStatus.ACCEPTED, self.manager.start(self._payload()))
                 return
@@ -742,6 +865,12 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help="调参 SQLite 与 trial 产物目录；默认位于训练工程 logs/auto_tuning",
     )
+    parser.add_argument(
+        "--pretrained-sources",
+        type=Path,
+        default=None,
+        help="可选旧管理员基础策略注册JSON；不配置也可直接上传单个.pt/.onnx",
+    )
     parser.add_argument(
         "--task", action="append", dest="tasks", help="允许前端启动的任务 ID；可重复"
     )
@@ -762,14 +891,23 @@ def main() -> None:
     if len(token) < 16:
         raise SystemExit("训练服务访问令牌至少需要 16 个字符")
     lease = GpuLease()
+    tuning_root = args.tuning_data_root or (Path(args.trainer_root) / "logs" / "auto_tuning")
+    sources = PretrainedSources(
+        args.pretrained_sources,
+        tuning_root / "pretrained_sources",
+        args.trainer_python,
+        args.trainer_root,
+    )
     manager = TrainingManager(
         args.trainer_root,
         args.trainer_python,
         tuple(args.tasks or DEFAULT_TASKS),
         lease=lease,
+        sources=sources,
+    )
+    tuning_manager = TuningManager(
+        args.trainer_root, args.trainer_python, tuning_root, lease, sources=sources
     )
-    tuning_root = args.tuning_data_root or (Path(args.trainer_root) / "logs" / "auto_tuning")
-    tuning_manager = TuningManager(args.trainer_root, args.trainer_python, tuning_root, lease)
     manager.preset_resolver = tuning_manager.preset_config
     TrainingRequestHandler.manager = manager
     TrainingRequestHandler.tuning_manager = tuning_manager