"""Persistent reward tuning session orchestrator.""" from __future__ import annotations import json import os import re import shutil import subprocess import threading from contextlib import suppress from copy import deepcopy from pathlib import Path from typing import Any from .advisor import AdvisorUnavailable, DeepSeekAdvisor from .process import GpuLease, ResourceBusyError, terminate_process from .schema import ( BASE_REWARD_CONFIGURATION, merge_proposal, validate_proposal, ) from .scoring import DEFAULT_OBJECTIVE_WEIGHTS, score_evaluation, validate_objective_weights from .storage import TuningStorage, now_iso from .study import OptunaStudies from .tensorboard import ingest_scalars ACTIVE_SESSION_STATES = { "queued", "running", "evaluating", "awaiting_approval", "paused", "interrupted", } RUNNING_STATES = {"queued", "running", "evaluating"} RUN_NAME = re.compile(r"^[A-Za-z0-9_.-]{1,64}$") class TuningError(RuntimeError): pass class TuningManager: def __init__( self, trainer_root: Path, python: str, data_root: Path, lease: GpuLease, advisor: Any | None = None, ): self.trainer_root = trainer_root.expanduser().resolve() self.python = python self.data_root = data_root.expanduser().resolve() self.data_root.mkdir(parents=True, exist_ok=True) self.storage = TuningStorage(self.data_root / "tuning.sqlite3") self.storage.recover_interrupted() self.studies = OptunaStudies(self.data_root) self.lease = lease self.advisor = advisor or DeepSeekAdvisor() self.lock = threading.RLock() self.condition = threading.Condition(self.lock) self.workers: dict[str, threading.Thread] = {} self.processes: dict[str, subprocess.Popen[str]] = {} self.cancel_events: dict[str, threading.Event] = {} def capability(self) -> dict[str, Any]: capability = self.advisor.capability() capability.update({"ready": (self.trainer_root / "scripts" / "evaluate.py").is_file()}) return capability @staticmethod def _integer(payload: dict, name: str, default: int, minimum: int, maximum: int) -> int: value = payload.get(name, default) if isinstance(value, bool) or not isinstance(value, int) or not minimum <= value <= maximum: raise TuningError(f"{name} 必须在 {minimum}–{maximum} 之间") return value def parse_create(self, payload: Any) -> tuple[str, dict, dict, bool]: if not isinstance(payload, dict): raise TuningError("请求体必须是 JSON 对象") mode = payload.get("mode", "automatic") if mode not in ("automatic", "approval"): raise TuningError("mode 必须是 automatic 或 approval") if payload.get("taskId", "Unitree-Go2-Flat") != "Unitree-Go2-Flat": raise TuningError("第一版只支持 Unitree-Go2-Flat") run_name = payload.get("runName", "auto-tune") if not isinstance(run_name, str) or not RUN_NAME.fullmatch(run_name): raise TuningError("runName 格式无效") gpu_ids = payload.get("gpuIds", [0]) if ( not isinstance(gpu_ids, list) or not gpu_ids or any( isinstance(value, bool) or not isinstance(value, int) or not 0 <= value <= 255 for value in gpu_ids ) ): raise TuningError("gpuIds 必须是非空非负整数数组") trial_count = self._integer(payload, "trialCount", 12, 4, 20) rung0 = self._integer(payload, "initialIterations", 300, 1, 1000000) rung1 = self._integer(payload, "middleIterations", 900, rung0, 1000000) rung2 = self._integer(payload, "finalIterations", 2000, rung1, 1000000) config = { "taskId": "Unitree-Go2-Flat", "numEnvs": self._integer(payload, "numEnvs", 4096, 1, 16384), "seed": self._integer(payload, "seed", 42, 0, 2147483647), "runName": run_name, "gpuIds": gpu_ids, "trialCount": trial_count, "rungs": [rung0, rung1, rung2], "promote": [trial_count, min(4, trial_count), min(2, trial_count)], "evalNumEnvs": self._integer(payload, "evalNumEnvs", 256, 1, 4096), "evalSteps": self._integer(payload, "evalSteps", 1000, 10, 100000), "earlyStopPatience": self._integer(payload, "earlyStopPatience", 4, 1, 20), } objective = validate_objective_weights( payload.get("objectiveWeights", DEFAULT_OBJECTIVE_WEIGHTS) ) fallback = payload.get("fallbackEnabled", False) if not isinstance(fallback, bool): raise TuningError("fallbackEnabled 必须是布尔值") return mode, config, objective, fallback def create(self, payload: Any) -> dict: mode, config, objective, fallback = self.parse_create(payload) capability = self.capability() if not capability["ready"]: raise TuningError("评估入口未就绪") if not capability["configured"] and not fallback: raise TuningError("DeepSeek Agent 未配置；设置 DEEPSEEK_API_KEY 或显式启用 fallback") with self.lock: active = [ session for session in self.storage.list_sessions() if session["state"] in ACTIVE_SESSION_STATES ] if active: raise ResourceBusyError("已有调参 session 未结束") session = self.storage.create_session(mode, config, objective, fallback) self._start_worker(session["id"], resume=False) return self.detail(session["id"]) def _start_worker(self, session_id: str, resume: bool) -> None: cancel = threading.Event() self.cancel_events[session_id] = cancel worker = threading.Thread( target=self._run_session, args=(session_id, resume, cancel), name=f"tuning-{session_id[:8]}", daemon=True, ) self.workers[session_id] = worker worker.start() def detail(self, session_id: str) -> dict: session = self.storage.get_session(session_id) session["trials"] = self.storage.list_trials(session_id) session["proposals"] = self.storage.list_proposals(session_id) session["audit"] = self.storage.audit_events(session_id) return session def list(self) -> list[dict]: return self.storage.list_sessions() def _session_root(self, session_id: str) -> Path: root = (self.data_root / "sessions" / session_id).resolve() if not root.is_relative_to(self.data_root): raise TuningError("非法 session 路径") root.mkdir(parents=True, exist_ok=True) return root def _run_command( self, session_id: str, command: list[str], cwd: Path, environment: dict[str, str], log_path: Path, ) -> int: owner = f"tuning:{session_id}" self.lease.acquire(owner) try: process = subprocess.Popen( command, cwd=cwd, env=environment, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, encoding="utf-8", errors="replace", bufsize=1, start_new_session=True, ) with self.lock: self.processes[session_id] = process assert process.stdout is not None with log_path.open("a", encoding="utf-8") as log: for line in process.stdout: log.write(line) log.flush() if self.cancel_events[session_id].is_set(): terminate_process(process) break return process.wait() finally: with self.lock: self.processes.pop(session_id, None) self.lease.release(owner) @staticmethod def _latest_checkpoint(run_dir: Path) -> Path | None: values = [] for path in run_dir.glob("model_*.pt"): match = re.fullmatch(r"model_(\d+)\.pt", path.name) if match: values.append((int(match.group(1)), path)) return max(values, default=(0, None), key=lambda value: value[0])[1] def _execute_trial( self, session: dict, trial: dict, resume_checkpoint: Path | None = None ) -> dict: session_id, trial_id = session["id"], trial["id"] root = self._session_root(session_id) run_dir = (root / trial["runDir"]).resolve() if not run_dir.is_relative_to(root): raise TuningError("trial 目录越界") run_dir.mkdir(parents=True, exist_ok=True) reward_path = run_dir / "reward_config.json" reward_path.write_text( json.dumps(trial["rewardConfig"], ensure_ascii=False, indent=2) + "\n", encoding="utf-8" ) config = session["config"] command = [ self.python, "-u", "scripts/train.py", config["taskId"], f"--env.scene.num-envs={config['numEnvs']}", f"--agent.max-iterations={trial['targetIterations']}", f"--agent.seed={config['seed']}", f"--agent.run-name={config['runName']}-t{trial['number']}-r{trial['rung']}", "--agent.logger=tensorboard", "--agent.upload-model=False", "--gpu-ids", json.dumps(config["gpuIds"], separators=(",", ":")), "--output-dir", str(run_dir), "--reward-config", str(reward_path), ] if resume_checkpoint is not None: command.extend(("--resume-checkpoint", str(resume_checkpoint))) environment = os.environ.copy() environment["WANDB_MODE"] = "disabled" environment["WANDB_SILENT"] = "true" self.storage.update_trial( trial_id, state="training", started_at=now_iso(), message="正在训练" ) self.storage.update_session( session_id, state="running", message=f"正在训练 trial {trial['number']} / rung {trial['rung']}", ) return_code = self._run_command( session_id, command, self.trainer_root, environment, run_dir / "train.log" ) ingest_scalars(self.storage, trial_id, run_dir) if self.cancel_events[session_id].is_set(): raise TuningError("session 已取消") checkpoint = self._latest_checkpoint(run_dir) policy = run_dir / "policy.onnx" if return_code != 0 or checkpoint is None or not policy.is_file(): raise TuningError(f"训练失败（返回码 {return_code}）或缺少 checkpoint/policy.onnx") self._wait_if_paused(session_id, self.cancel_events[session_id]) eval_output = run_dir / "evaluation.json" eval_command = [ self.python, "-u", "scripts/evaluate.py", config["taskId"], "--checkpoint", str(checkpoint), "--output", str(eval_output), "--reward-config", str(reward_path), f"--num-envs={config['evalNumEnvs']}", f"--steps-per-seed={config['evalSteps']}", "--gpu-ids", json.dumps(config["gpuIds"], separators=(",", ":")), ] self.storage.update_trial(trial_id, state="evaluating", message="正在固定协议评估") self.storage.update_session( session_id, state="evaluating", message=f"正在评估 trial {trial['number']}" ) return_code = self._run_command( session_id, eval_command, self.trainer_root, environment, run_dir / "evaluate.log" ) ingest_scalars(self.storage, trial_id, run_dir / "evaluation-events") if return_code != 0 or not eval_output.is_file(): raise TuningError(f"评估失败（返回码 {return_code}）") evaluation = json.loads(eval_output.read_text(encoding="utf-8")) baseline_trial = self.storage.list_trials(session_id)[0] if baseline_trial["evaluation"] is None: scored = { "eligible": True, "score": 0.0, "components": {}, "metrics": evaluation["metrics"], } else: scored = score_evaluation( baseline_trial["evaluation"]["metrics"], evaluation["metrics"], session["objectiveWeights"], ) evaluation["score"] = scored rel_checkpoint = str(checkpoint.relative_to(root)) rel_policy = str(policy.relative_to(root)) self.storage.update_trial( trial_id, state="completed", ended_at=now_iso(), message="训练与评估完成", checkpoint_path=rel_checkpoint, policy_path=rel_policy, evaluation=evaluation, score=scored["score"], eligible=scored["eligible"], ) try: optuna_number = self.studies.record( session_id, trial["rewardConfig"], scored["score"], scored["eligible"], trial["rung"], ) self.storage.audit( session_id, "optuna_trial_recorded", {"trialId": trial_id, "number": optuna_number} ) except Exception as error: self.storage.audit( session_id, "optuna_record_failed", {"trialId": trial_id, "error": str(error)[:500]} ) return self.storage.get_trial(trial_id) def _best(self, session_id: str, rung: int | None = None) -> dict | None: trials = [ trial for trial in self.storage.list_trials(session_id) if trial["state"] == "completed" and trial["eligible"] ] if rung is not None: trials = [trial for trial in trials if trial["rung"] == rung] return max( trials, key=lambda trial: trial["score"] if trial["score"] is not None else -999, default=None, ) def _proposal_context(self, session: dict) -> dict: trials = self.storage.list_trials(session["id"])[-12:] rejected_feedback = [ proposal["feedback"] for proposal in self.storage.list_proposals(session["id"]) if proposal["state"] == "rejected" and proposal["feedback"] ][-4:] return { "task": session["config"]["taskId"], "objectiveWeights": session["objectiveWeights"], "allowlist": "服务端将验证固定 schema；最多四项修改", "rejectedFeedback": rejected_feedback, "trials": [ { "number": t["number"], "rung": t["rung"], "score": t["score"], "eligible": t["eligible"], "rewardConfig": t["rewardConfig"], "evaluation": t["evaluation"] and t["evaluation"].get("metrics"), } for t in trials ], } def _fallback_patch(self, previous: dict, index: int) -> dict: names = ("track_linear_velocity", "action_rate_l2", "body_orientation_l2", "foot_slip") name = names[index % len(names)] old = previous["weights"][name] factor = 1.1 if index % 2 == 0 else 0.9 return validate_proposal({"weights": {name: old * factor}}, previous) def _request_proposal( self, session: dict, previous: dict, base_trial_id: str, index: int ) -> dict: try: result = self.advisor.propose(self._proposal_context(session), previous) source = "agent" except Exception as error: if not session["fallbackEnabled"]: raise AdvisorUnavailable(str(error)) from error result = { "patch": self._fallback_patch(previous, index), "rationale": f"Agent 不可用，显式 fallback：{error}", "expectedImpact": {}, "confidence": 0.2, "promptHash": None, "usage": {}, "model": "optuna-fallback", } source = "fallback" proposal = self.storage.create_proposal( session["id"], base_trial_id, result["patch"], result["rationale"], result.get("expectedImpact", {}), result["confidence"], source, ) self.storage.audit( session["id"], "proposal_created", { "proposalId": proposal["id"], "source": source, "promptHash": result.get("promptHash"), "usage": result.get("usage", {}), "model": result.get("model"), }, ) return proposal def _wait_for_approval( self, session_id: str, proposal_id: str, cancel: threading.Event ) -> dict: with self.condition: while not cancel.is_set(): proposal = self.storage.get_proposal(proposal_id) if proposal["state"] != "pending": return proposal self.condition.wait(timeout=1.0) raise TuningError("session 已取消") def _run_session(self, session_id: str, resume: bool, cancel: threading.Event) -> None: try: session = self.storage.get_session(session_id) trials = self.storage.list_trials(session_id) if resume: root = self._session_root(session_id) for interrupted in [trial for trial in trials if trial["state"] == "interrupted"]: run_dir = (root / interrupted["runDir"]).resolve() if run_dir.is_relative_to(root): shutil.rmtree(run_dir, ignore_errors=True) self.storage.delete_trial(interrupted["id"]) self.storage.audit( session_id, "session_resumed", { "discardedInterruptedTrials": [ t["id"] for t in trials if t["state"] == "interrupted" ] }, ) trials = self.storage.list_trials(session_id) if not trials: baseline_dir = "trial-000-rung-0" trial = self.storage.create_trial( session_id, 0, 0, session["config"]["rungs"][0], deepcopy(BASE_REWARD_CONFIGURATION), None, baseline_dir, ) self._execute_trial(session, trial) session = self.storage.get_session(session_id) completed_rung0 = [ t for t in self.storage.list_trials(session_id) if t["rung"] == 0 and t["state"] == "completed" ] next_number = len({t["number"] for t in completed_rung0}) best_score = max((t["score"] or 0.0 for t in completed_rung0), default=0.0) no_improve = 0 while ( next_number < session["config"]["trialCount"] and no_improve < session["config"]["earlyStopPatience"] ): if cancel.is_set(): raise TuningError("session 已取消") self._wait_if_paused(session_id, cancel) base = self._best(session_id, rung=0) or completed_rung0[0] proposal = self._request_proposal( session, base["rewardConfig"], base["id"], next_number ) if session["mode"] == "approval": self.storage.update_session( session_id, state="awaiting_approval", message="等待批准 Agent 建议" ) proposal = self._wait_for_approval(session_id, proposal["id"], cancel) if proposal["state"] == "rejected": self.storage.audit( session_id, "proposal_rejected", {"proposalId": proposal["id"], "feedback": proposal["feedback"]}, ) continue else: self.storage.decide_proposal(proposal["id"], "approved", "自动模式") proposal = self.storage.get_proposal(proposal["id"]) self._wait_if_paused(session_id, cancel) reward_config = merge_proposal(base["rewardConfig"], proposal["patch"]) trial = self.storage.create_trial( session_id, next_number, 0, session["config"]["rungs"][0], reward_config, proposal["id"], f"trial-{next_number:03d}-rung-0", ) result = self._execute_trial(session, trial) if result["eligible"] and (result["score"] or -999) > best_score + 0.01: best_score = result["score"] no_improve = 0 else: no_improve += 1 self.storage.update_session(session_id, consecutive_no_improve=no_improve) next_number += 1 # Promote top configurations; each new rung resumes its own previous checkpoint. for rung in (1, 2): self._wait_if_paused(session_id, cancel) previous = [ t for t in self.storage.list_trials(session_id) if t["rung"] == rung - 1 and t["state"] == "completed" and t["eligible"] ] previous.sort(key=lambda t: t["score"] or -999, reverse=True) promoted_numbers = { trial["number"] for trial in self.storage.list_trials(session_id) if trial["rung"] == rung and trial["state"] == "completed" } for parent in previous[: session["config"]["promote"][rung]]: if parent["number"] in promoted_numbers: continue if cancel.is_set(): raise TuningError("session 已取消") root = self._session_root(session_id) checkpoint = root / parent["checkpointPath"] trial = self.storage.create_trial( session_id, parent["number"], rung, session["config"]["rungs"][rung], parent["rewardConfig"], parent["proposalId"], f"trial-{parent['number']:03d}-rung-{rung}", ) self._execute_trial(session, trial, checkpoint) best = ( self._best(session_id, rung=2) or self._best(session_id, rung=1) or self._best(session_id, rung=0) ) if best is None: raise TuningError("没有通过安全门槛的 trial") preset_name = f"{session['config']['runName']}-{session_id[:8]}" self.storage.save_preset(preset_name, session_id, best["id"], best["rewardConfig"]) self.storage.update_session( session_id, state="succeeded", best_trial_id=best["id"], current_trial_id=None, message="调参完成", ) self.storage.audit( session_id, "session_completed", {"bestTrialId": best["id"], "preset": preset_name} ) except Exception as error: state = self.storage.get_session(session_id)["state"] if cancel.is_set() or state == "cancelled": self.storage.update_session( session_id, state="cancelled", message="调参已取消", current_trial_id=None ) else: self.storage.update_session( session_id, state="failed", message=str(error), current_trial_id=None ) self.storage.audit(session_id, "session_failed", {"error": str(error)[:1000]}) finally: with self.lock: self.workers.pop(session_id, None) self.processes.pop(session_id, None) def _wait_if_paused(self, session_id: str, cancel: threading.Event) -> None: with self.condition: while self.storage.get_session(session_id)["state"] == "paused" and not cancel.is_set(): self.condition.wait(timeout=1.0) def approve(self, session_id: str, proposal_id: str, payload: Any) -> dict: proposal = self.storage.get_proposal(proposal_id) if proposal["sessionId"] != session_id: raise TuningError("proposal 不属于该 session") patch = proposal["patch"] feedback = None if isinstance(payload, dict): feedback = payload.get("feedback") if "patch" in payload: base = self.storage.get_trial(proposal["baseTrialId"]) patch = validate_proposal(payload["patch"], base["rewardConfig"]) if not self.storage.decide_proposal(proposal_id, "approved", feedback, patch): raise TuningError("proposal 已处理") self.storage.audit( session_id, "proposal_approved", {"proposalId": proposal_id, "modified": patch != proposal["patch"]}, ) with self.condition: self.condition.notify_all() return self.detail(session_id) def reject(self, session_id: str, proposal_id: str, payload: Any) -> dict: feedback = payload.get("feedback", "") if isinstance(payload, dict) else "" if not isinstance(feedback, str) or len(feedback) > 2000: raise TuningError("feedback 无效") proposal = self.storage.get_proposal(proposal_id) if proposal["sessionId"] != session_id: raise TuningError("proposal 不属于该 session") if not self.storage.decide_proposal(proposal_id, "rejected", feedback): raise TuningError("proposal 已处理") with self.condition: self.condition.notify_all() return self.detail(session_id) def pause(self, session_id: str) -> dict: session = self.storage.get_session(session_id) if session["state"] not in RUNNING_STATES | {"awaiting_approval"}: raise TuningError("当前状态不能暂停") self.storage.update_session( session_id, state="paused", message="已暂停后续调度；当前子进程将完成" ) return self.detail(session_id) def resume(self, session_id: str) -> dict: session = self.storage.get_session(session_id) if session["state"] == "paused": self.storage.update_session(session_id, state="running", message="继续调参") with self.condition: self.condition.notify_all() elif session["state"] == "interrupted": self.storage.update_session(session_id, state="queued", message="从最近完整结果恢复") self._start_worker(session_id, resume=True) else: raise TuningError("当前状态不能恢复") return self.detail(session_id) def cancel(self, session_id: str) -> dict: self.storage.get_session(session_id) self.storage.update_session(session_id, state="cancelled", message="正在取消") event = self.cancel_events.get(session_id) if event: event.set() process = self.processes.get(session_id) if process: terminate_process(process) with self.condition: self.condition.notify_all() return self.detail(session_id) def metrics( self, session_id: str, trial_id: str, tags: list[str] | None, max_points: int ) -> dict: trial = self.storage.get_trial(trial_id) if trial["sessionId"] != session_id: raise TuningError("trial 不属于该 session") return {"trialId": trial_id, "series": self.storage.metrics(trial_id, tags, max_points)} def best_artifact(self, session_id: str) -> Path: session = self.storage.get_session(session_id) if not session["bestTrialId"]: raise TuningError("尚无最佳策略") trial = self.storage.get_trial(session["bestTrialId"]) if not trial["policyPath"]: raise TuningError("最佳策略文件不存在") root = self._session_root(session_id) path = (root / trial["policyPath"]).resolve() if not path.is_relative_to(root) or not path.is_file(): raise TuningError("最佳策略文件不存在") return path def preset_config(self, preset_id: str) -> dict: return self.storage.get_preset(preset_id)["rewardConfig"] def test_agent(self) -> dict: try: return self.advisor.test_connection() except Exception as error: raise TuningError(f"Agent 连接测试失败：{error}") from error def shutdown(self) -> None: for session_id in list(self.workers): with suppress(KeyError, TuningError): self.cancel(session_id) for worker in list(self.workers.values()): worker.join(timeout=7)