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

2026-09-08 10:50:13 +08:00
parent fa5485049a
commit 438e56bcc8
113 changed files with 15027 additions and 539 deletions
@@ -0,0 +1,447 @@
+"""Administrator-registered, content-bound local training sources (no client paths)."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import re
+import shutil
+import stat
+import subprocess
+import tempfile
+import threading
+import time
+from contextlib import contextmanager
+from copy import deepcopy
+from pathlib import Path
+
+SOURCE_ID = re.compile(r"^[A-Za-z0-9_-]{1,64}$")
+DIGEST = re.compile(r"^[0-9a-f]{64}$")
+LIMITS = {
+    "checkpoint": 256 * 1024**2,
+    "onnx": 64 * 1024**2,
+    "env": 2 * 1024**2,
+    "agent": 128 * 1024,
+}
+TASKS = ["Unitree-Go2-Flat", "Unitree-Go2-ObstacleAvoidance"]
+
+
+class SourceError(ValueError):
+    pass
+
+
+def regular_bytes(path: Path, root: Path, limit: int) -> bytes:
+    """Walk with dir_fd/O_NOFOLLOW, so symlink substitution cannot escape the root."""
+    if ".." in path.parts or not path.is_absolute():
+        raise SourceError("基础策略路径必须是允许根内的绝对路径，不能含 ..")
+    try:
+        parts = path.relative_to(root).parts
+    except ValueError as error:
+        raise SourceError("基础策略文件不在管理员配置的允许根内") from error
+    if not parts:
+        raise SourceError("基础策略必须是常规文件")
+    try:
+        fd = os.open(root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
+    except OSError as error:
+        raise SourceError("基础策略快照目录缺失或无效，拒绝训练") from error
+    try:
+        for part in parts[:-1]:
+            next_fd = os.open(part, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=fd)
+            os.close(fd)
+            fd = next_fd
+        leaf = os.open(parts[-1], os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=fd)
+        with os.fdopen(leaf, "rb") as stream:
+            info = os.fstat(stream.fileno())
+            if not stat.S_ISREG(info.st_mode) or info.st_size > limit:
+                raise SourceError("基础策略文件必须是常规文件且不能超过大小上限")
+            data = stream.read(limit + 1)
+            if len(data) > limit:
+                raise SourceError("基础策略文件超过大小上限")
+            return data
+    except OSError as error:
+        raise SourceError(
+            "基础策略文件缺失或含symlink；请提供.pt、policy.onnx和params配置常规文件"
+        ) from error
+    finally:
+        os.close(fd)
+
+
+class PretrainedSources:
+    def __init__(self, config_path: Path | None, store: Path, python: str, trainer_root: Path):
+        self.store = store.expanduser().resolve()
+        self.store.mkdir(parents=True, exist_ok=True, mode=0o700)
+        self.python = python
+        self.trainer_root = trainer_root.resolve()
+        self.entries: dict[str, dict] = {}
+        self._upload_lock = threading.Lock()
+        self._catalog_lock = threading.RLock()
+        self._restore_uploads()
+        if config_path is None:
+            return
+        try:
+            if config_path.stat().st_size > 128 * 1024:
+                raise SourceError("基础策略注册配置不能超过128KiB")
+            config = json.loads(config_path.read_text())
+            if not isinstance(config, dict) or set(config) != {"allowedRoots", "sources"}:
+                raise SourceError("注册配置只接受allowedRoots和sources")
+            if (
+                not isinstance(config["allowedRoots"], list)
+                or not isinstance(config["sources"], list)
+                or len(config["sources"]) > 32
+            ):
+                raise SourceError("允许根/sources必须为数组，最多32个注册条目")
+            roots = [Path(p).expanduser().resolve(strict=True) for p in config["allowedRoots"]]
+            if not roots or any(not p.is_dir() for p in roots):
+                raise SourceError("请配置存在的本地允许根目录")
+            for entry in config["sources"]:
+                if not isinstance(entry, dict) or set(entry) != {
+                    "id",
+                    "label",
+                    "checkpoint",
+                    "onnx",
+                }:
+                    raise SourceError("注册条目必须包含id/label/checkpoint/onnx")
+                key = entry["id"]
+                if not isinstance(key, str) or not SOURCE_ID.fullmatch(key) or key in self.entries:
+                    raise SourceError("注册id无效或重复")
+                if not isinstance(entry["label"], str) or not 1 <= len(entry["label"]) <= 100:
+                    raise SourceError("基础策略label必须为1–100字符")
+                record = {"id": key, "label": entry["label"], "ready": False, "compatibleTasks": []}
+                self.entries[key] = {"public": record}
+                try:
+                    checkpoint, onnx = (
+                        Path(entry["checkpoint"]).expanduser(),
+                        Path(entry["onnx"]).expanduser(),
+                    )
+                    if not re.fullmatch(r"[A-Za-z0-9_.-]+\.pt", checkpoint.name):
+                        raise SourceError("请选择ONNX对应的.pt训练checkpoint，ONNX不能直接续训")
+                    files = {
+                        "checkpoint": checkpoint,
+                        "onnx": onnx,
+                        "env": checkpoint.parent / "params/env.yaml",
+                        "agent": checkpoint.parent / "params/agent.yaml",
+                    }
+                    data = {}
+                    for name, path in files.items():
+                        root = next((r for r in roots if path.is_relative_to(r)), None)
+                        if root is None:
+                            raise SourceError("注册文件不在允许根内")
+                        data[name] = regular_bytes(path, root, LIMITS[name])
+                    names = {
+                        "checkpoint": checkpoint.name,
+                        "onnx": "policy.onnx",
+                        "env": "params/env.yaml",
+                        "agent": "params/agent.yaml",
+                    }
+                    with tempfile.TemporaryDirectory(dir=self.store, prefix="import-") as temporary:
+                        directory = Path(temporary)
+                        for name, content in data.items():
+                            destination = directory / names[name]
+                            destination.parent.mkdir(exist_ok=True)
+                            destination.write_bytes(content)
+                        manifest = self._validate(directory, checkpoint.name, TASKS[0], None)
+                        digest = manifest["source_id"]
+                        bound = {
+                            "sourceId": digest,
+                            "registeredId": key,
+                            "label": entry["label"],
+                            "manifest": manifest,
+                        }
+                        destination = self.store / digest
+                        if not destination.exists():
+                            shutil.copytree(directory, destination)
+                            for file in destination.rglob("*"):
+                                file.chmod(0o555 if file.is_dir() else 0o444)
+                            destination.chmod(0o555)
+                        self.verify(bound)
+                    record.update(
+                        id=digest,
+                        ready=True,
+                        compatibleTasks=TASKS,
+                        observationSizes=[47, 81, 97],
+                        initialization=bound,
+                    )
+                except (SourceError, OSError, ValueError, TypeError) as error:
+                    record["error"] = f"{error}；请修正管理员注册配置并重启服务，不会退回随机初始化"
+        except (OSError, TypeError, ValueError) as error:
+            raise SourceError(f"基础策略注册配置无效：{error}") from error
+
+    def _restore_uploads(self):
+        # Only committed content directories are catalogued. Partial uploads are never sources.
+        for directory in self.store.iterdir():
+            if not DIGEST.fullmatch(directory.name) or not (directory / "upload.json").is_file():
+                continue
+            try:
+                manifest = json.loads(
+                    regular_bytes(directory / "upload.json", self.store, 64 * 1024)
+                )
+                label = json.loads(regular_bytes(directory / "label.json", self.store, 1024))[
+                    "label"
+                ]
+                bound = {
+                    "sourceId": directory.name,
+                    "registeredId": directory.name,
+                    "label": self._display_name(label),
+                    "manifest": manifest,
+                }
+                if manifest["source_id"] != directory.name:
+                    raise SourceError("上传来源身份损坏")
+                self.verify(bound)
+                self._publish_upload(bound)
+            except (SourceError, ValueError, KeyError, TypeError, OSError):
+                self.entries[directory.name] = {
+                    "public": {
+                        "id": directory.name,
+                        "label": "失效的已上传策略",
+                        "ready": False,
+                        "compatibleTasks": [],
+                        "error": "上传快照损坏，拒绝随机初始化；请重新上传原文件或恢复快照",
+                    }
+                }
+
+    @staticmethod
+    def _display_name(name):
+        if not isinstance(name, str):
+            return "uploaded-policy"
+        name = name.replace("\\", "/").split("/")[-1]
+        return (
+            re.sub(r"[^\w .()\-]", "_", name, flags=re.UNICODE)[:100].strip(" .")
+            or "uploaded-policy"
+        )
+
+    def _publish_upload(self, bound):
+        record = {
+            "id": bound["sourceId"],
+            "label": bound["label"],
+            "ready": True,
+            "compatibleTasks": TASKS,
+            "observationSizes": [47, 81, 97],
+            "initialization": bound,
+        }
+        with self._catalog_lock:
+            self.entries[bound["sourceId"]] = {"public": record}
+        return deepcopy(record)
+
+    @contextmanager
+    def upload_slot(self, length):
+        # At most one receiving/decoding upload, 32 sources and 2GiB cumulative store.
+        if not self._upload_lock.acquire(blocking=False):
+            raise SourceError("已有文件正在上传/验证，请稍后重试")
+        try:
+            files = list(self.store.rglob("*"))
+            used = sum(p.lstat().st_size for p in files if not p.is_symlink() and p.is_file())
+            count = sum(1 for p in self.store.iterdir() if DIGEST.fullmatch(p.name))
+            if count >= 32 or used + length + 16 * 1024**2 > 2 * 1024**3:
+                raise SourceError("基础策略存储已达32个来源或2GiB上限，请由本机维护者释放空间")
+            with tempfile.TemporaryDirectory(dir=self.store, prefix="upload-") as temporary:
+                yield Path(temporary)
+        finally:
+            self._upload_lock.release()
+
+    def receive_upload(self, stream, length, fmt, template, display_name, *, set_timeout=None):
+        if template != "go2-legacy47-v1":
+            raise SourceError("必须明确确认go2-legacy47-v1观测与关节模板")
+        limit = {"pt": 256 * 1024**2, "onnx": 64 * 1024**2}.get(fmt)
+        if limit is None:
+            raise SourceError("仅支持单个.pt或.onnx，不接受ZIP/路径/配套目录")
+        if type(length) is not int or not 0 < length <= limit:
+            raise SourceError("上传文件为空或超过.pt 256MiB / ONNX 64MiB上限")
+        with self.upload_slot(length) as directory:
+            path = directory / f"upload.{fmt}"
+            deadline = time.monotonic() + 60
+            remaining = length
+            try:
+                with path.open("xb") as output:
+                    while remaining:
+                        timeout = deadline - time.monotonic()
+                        if timeout <= 0:
+                            raise SourceError("上传超时，请重试")
+                        try:
+                            if set_timeout is not None:
+                                set_timeout(min(10, timeout))
+                            # BufferedReader.read(n) may perform many recv calls whose socket
+                            # timeouts reset with each trickled byte. read1 returns after one
+                            # raw read, allowing the total deadline to be checked every time.
+                            read_chunk = getattr(stream, "read1", stream.read)
+                            chunk = read_chunk(min(1024 * 1024, remaining))
+                        except OSError as error:
+                            raise SourceError("上传超时/连接中断，临时文件已清理") from error
+                        if not chunk or len(chunk) > remaining:
+                            raise SourceError("上传连接中断或实际长度不符")
+                        output.write(chunk)
+                        remaining -= len(chunk)
+            finally:
+                if set_timeout is not None:
+                    set_timeout(10)
+            request = {
+                "path": str(path),
+                "directory": str(directory),
+                "format": fmt,
+                "template": template,
+            }
+            env = {
+                **os.environ,
+                "CUDA_VISIBLE_DEVICES": "",
+                "OMP_NUM_THREADS": "1",
+                "OPENBLAS_NUM_THREADS": "1",
+                "MKL_NUM_THREADS": "1",
+            }
+            try:
+                result = subprocess.run(
+                    [self.python, str(Path(__file__).parent / "rl/scripts/validate_upload.py")],
+                    input=json.dumps(request),
+                    capture_output=True,
+                    text=True,
+                    cwd=self.trainer_root,
+                    env=env,
+                    timeout=60,
+                )
+                manifest = json.loads(result.stdout)
+            except (OSError, subprocess.TimeoutExpired, ValueError) as error:
+                raise SourceError("模型验证超时/资源超限或训练Python依赖不可用") from error
+            if result.returncode:
+                raise SourceError(manifest.get("error", "不支持的模型文件"))
+            digest = manifest["source_id"]
+            if not DIGEST.fullmatch(digest):
+                raise SourceError("模型验证返回无效身份")
+            label = self._display_name(display_name)
+            (directory / "upload.json").write_text(json.dumps(manifest), encoding="utf-8")
+            (directory / "label.json").write_text(json.dumps({"label": label}), encoding="utf-8")
+            destination = self.store / digest
+            if destination.exists():
+                # Dedup never replaces a committed artifact or changes an old job binding.
+                old = next((r for r in self.catalog() if r["id"] == digest and r["ready"]), None)
+                if old is None:
+                    raise SourceError("同ID旧快照已损坏，请恢复快照后重试；不会覆盖旧任务来源")
+                self.verify(old["initialization"])
+                return old
+            for file in directory.iterdir():
+                file.chmod(0o444)
+            directory.rename(destination)
+            destination.chmod(0o555)
+            bound = {
+                "sourceId": digest,
+                "registeredId": digest,
+                "label": label,
+                "manifest": manifest,
+            }
+            self.verify(bound)
+            return self._publish_upload(bound)
+
+    def _validate(self, directory: Path, checkpoint: str, task_id: str, task_config: dict | None):
+        command = [self.python, str(Path(__file__).parent / "rl/scripts/validate_pretrained.py")]
+        request = {
+            "directory": str(directory),
+            "checkpoint": checkpoint,
+            "taskId": task_id,
+            "taskConfig": task_config,
+            "uploaded": (directory / "upload.json").is_file(),
+        }
+        try:
+            result = subprocess.run(
+                command,
+                input=json.dumps(request),
+                cwd=self.trainer_root,
+                capture_output=True,
+                text=True,
+                timeout=60,
+            )
+        except (subprocess.TimeoutExpired, OSError) as error:
+            raise SourceError("基础策略验证器超时或不可用，请检查训练Python环境") from error
+        if result.returncode:
+            # Validator emits only an actionable error, never a traceback or file contents.
+            raise SourceError(
+                result.stderr.strip()[-1500:] or "基础策略验证器失败，请检查训练Python依赖"
+            )
+        try:
+            return json.loads(result.stdout)
+        except ValueError as error:
+            raise SourceError("基础策略验证器未返回有效身份描述") from error
+
+    def catalog(self):
+        with self._catalog_lock:
+            sources = {}
+            for entry in self.entries.values():
+                sources.setdefault(entry["public"]["id"], deepcopy(entry["public"]))
+            return list(sources.values())
+
+    def bind(self, source_id, task_id: str, task_config=None):
+        # Selection is content-addressed, so a stale UI cannot silently bind a newly
+        # registered model under the same administrator-friendly registration name.
+        entry = next(
+            (entry for entry in self.catalog() if entry["id"] == source_id),
+            None,
+        )
+        if not isinstance(source_id, str) or entry is None:
+            raise SourceError("基础策略内容ID不存在或已变化；请刷新服务连接，不接受文件路径")
+        if not entry["ready"]:
+            raise SourceError(entry["error"])
+        if task_id not in entry["compatibleTasks"]:
+            raise SourceError("该基础策略仅兼容Flat47/Obstacle81或97；Rough高度扫描不支持迁移")
+        bound = deepcopy(entry["initialization"])
+        directory = self.verify(bound)
+        manifest = self._validate(
+            directory, bound["manifest"]["artifacts"]["checkpoint"]["name"], task_id, task_config
+        )
+        if manifest["source_id"] != bound["sourceId"]:
+            raise SourceError("基础策略快照身份变化，拒绝初始化")
+        return bound
+
+    def verify(self, bound: dict) -> Path:
+        try:
+            digest = bound["sourceId"]
+            if not isinstance(digest, str) or not DIGEST.fullmatch(digest):
+                raise SourceError("基础策略快照身份无效")
+            directory = self.store / digest
+            artifacts = bound["manifest"]["artifacts"]
+            if "sourceFormat" in bound["manifest"]:
+                fmt = bound["manifest"]["sourceFormat"]
+                if fmt not in ("pt", "onnx"):
+                    raise SourceError("上传格式无效")
+                stored = json.loads(regular_bytes(directory / "upload.json", self.store, 64 * 1024))
+                if stored != bound["manifest"]:
+                    raise SourceError("上传manifest SHA绑定变化，拒绝初始化")
+                for name, relative, limit in (
+                    (
+                        "upload",
+                        f"upload.{fmt}",
+                        LIMITS["checkpoint"] if fmt == "pt" else LIMITS["onnx"],
+                    ),
+                    ("checkpoint", "actor.pt", 16 * 1024**2),
+                ):
+                    content = regular_bytes(directory / relative, self.store, limit)
+                    if (
+                        artifacts[name]["name"] != relative
+                        or hashlib.sha256(content).hexdigest() != artifacts[name]["sha256"]
+                    ):
+                        raise SourceError("上传快照SHA不匹配，拒绝训练")
+                return directory
+            for name, relative in {
+                "checkpoint": artifacts["checkpoint"]["name"],
+                "onnx": "policy.onnx",
+                "env": "params/env.yaml",
+                "agent": "params/agent.yaml",
+            }.items():
+                content = regular_bytes(directory / relative, self.store, LIMITS[name])
+                if hashlib.sha256(content).hexdigest() != artifacts[name]["sha256"]:
+                    raise SourceError("基础策略快照SHA不匹配，拒绝训练；请恢复原快照")
+            return directory
+        except (KeyError, TypeError) as error:
+            raise SourceError("持久化基础策略描述损坏，拒绝训练") from error
+
+    def arguments(self, bound: dict) -> list[str]:
+        directory = self.verify(bound)
+        upload_args = (
+            ["--pretrained-upload-manifest", str(directory / "upload.json")]
+            if "sourceFormat" in bound["manifest"]
+            else []
+        )
+        return upload_args + [
+            "--pretrained-checkpoint",
+            str(directory / bound["manifest"]["artifacts"]["checkpoint"]["name"]),
+            "--pretrained-allowed-roots",
+            json.dumps([str(directory)]),
+            "--pretrained-source-id",
+            bound["sourceId"],
+        ]