138 lines
5.5 KiB
Python
138 lines
5.5 KiB
Python
"""PydanticAI adapter for the DeepSeek reward-tuning advisor."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from .schema import validate_proposal
|
|
|
|
SYSTEM_PROMPT = """你是 Unitree Go2 强化学习奖励调参专家。
|
|
只根据提供的数值配置、训练曲线摘要和固定评估结果提出下一轮稀疏修改。
|
|
必须优先保持当前任务的客观安全门槛：Flat速度跟踪/跌倒，Obstacle无碰撞到达/跌倒。
|
|
每轮最多修改四个白名单标量，不得改变符号、函数、传感器、评估协议或结构。
|
|
不要建议 Python 代码、命令、文件路径或白名单外参数。输出必须符合 RewardProposal schema。
|
|
"""
|
|
|
|
|
|
class AdvisorUnavailable(RuntimeError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AdvisorConfig:
|
|
api_key: str | None
|
|
base_url: str = "https://api.deepseek.com"
|
|
model: str = "deepseek-v4-flash"
|
|
|
|
@classmethod
|
|
def from_environment(cls) -> AdvisorConfig:
|
|
return cls(
|
|
api_key=os.environ.get("DEEPSEEK_API_KEY"),
|
|
base_url=os.environ.get("MUJOCO_TUNING_AGENT_BASE_URL", "https://api.deepseek.com"),
|
|
model=os.environ.get("MUJOCO_TUNING_AGENT_MODEL", "deepseek-v4-flash"),
|
|
)
|
|
|
|
|
|
class DeepSeekAdvisor:
|
|
def __init__(self, config: AdvisorConfig | None = None):
|
|
self.config = config or AdvisorConfig.from_environment()
|
|
self._cached_agent = None
|
|
|
|
def capability(self) -> dict[str, Any]:
|
|
try:
|
|
import pydantic_ai # noqa: F401
|
|
except ImportError:
|
|
installed = False
|
|
else:
|
|
installed = True
|
|
return {
|
|
"configured": bool(self.config.api_key) and installed,
|
|
"apiKeyConfigured": bool(self.config.api_key),
|
|
"frameworkInstalled": installed,
|
|
"model": self.config.model,
|
|
"baseUrl": self.config.base_url,
|
|
}
|
|
|
|
def _agent(self):
|
|
if self._cached_agent is not None:
|
|
return self._cached_agent
|
|
if not self.config.api_key:
|
|
raise AdvisorUnavailable("未配置 DEEPSEEK_API_KEY")
|
|
try:
|
|
import httpx2
|
|
from pydantic import BaseModel, Field
|
|
from pydantic_ai import Agent, PromptedOutput
|
|
from pydantic_ai.models.openai import OpenAIChatModel
|
|
from pydantic_ai.providers.openai import OpenAIProvider
|
|
except ImportError as error:
|
|
raise AdvisorUnavailable(
|
|
"缺少 PydanticAI，请安装 training_server/requirements.txt"
|
|
) from error
|
|
|
|
class RewardProposalOutput(BaseModel):
|
|
weights: dict[str, float] = Field(default_factory=dict)
|
|
params: dict[str, float] = Field(default_factory=dict)
|
|
rationale: str = Field(min_length=1, max_length=2000)
|
|
expected_impact: dict[str, str] = Field(default_factory=dict)
|
|
confidence: float = Field(ge=0.0, le=1.0)
|
|
|
|
proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("ALL_PROXY")
|
|
if proxy and proxy.startswith("socks://"):
|
|
proxy = "socks5://" + proxy.removeprefix("socks://")
|
|
http_client = httpx2.AsyncClient(proxy=proxy, trust_env=False, timeout=60.0)
|
|
provider = OpenAIProvider(
|
|
base_url=self.config.base_url, api_key=self.config.api_key, http_client=http_client
|
|
)
|
|
model = OpenAIChatModel(self.config.model, provider=provider) # type: ignore[arg-type]
|
|
self._cached_agent = Agent(
|
|
model,
|
|
output_type=PromptedOutput(RewardProposalOutput),
|
|
system_prompt=SYSTEM_PROMPT,
|
|
retries=2,
|
|
model_settings={"temperature": 0.2},
|
|
)
|
|
return self._cached_agent
|
|
|
|
def propose(self, context: dict[str, Any], previous: dict) -> dict[str, Any]:
|
|
prompt = json.dumps(context, ensure_ascii=False, separators=(",", ":"), allow_nan=False)
|
|
result = self._agent().run_sync(prompt)
|
|
output = result.output
|
|
patch = validate_proposal(
|
|
{"weights": dict(output.weights), "params": dict(output.params)},
|
|
previous,
|
|
task_id=context.get("task", "Unitree-Go2-Flat"),
|
|
)
|
|
try:
|
|
usage = result.usage()
|
|
usage_value = {
|
|
key: getattr(usage, key)
|
|
for key in ("requests", "input_tokens", "output_tokens", "total_tokens")
|
|
if getattr(usage, key, None) is not None
|
|
}
|
|
except (AttributeError, TypeError):
|
|
usage_value = {}
|
|
return {
|
|
"patch": patch,
|
|
"rationale": output.rationale,
|
|
"expectedImpact": dict(output.expected_impact),
|
|
"confidence": float(output.confidence),
|
|
"promptHash": hashlib.sha256(prompt.encode()).hexdigest(),
|
|
"usage": usage_value,
|
|
"model": self.config.model,
|
|
}
|
|
|
|
def test_connection(self) -> dict[str, Any]:
|
|
base = {
|
|
"weights": {"track_linear_velocity": 1.0},
|
|
"params": {},
|
|
"instruction": "仅返回一个合法示例：把 track_linear_velocity 改为 1.1。",
|
|
}
|
|
# A minimal full previous config is supplied by callers for actual proposals;
|
|
# connectivity probing only verifies the provider and structured response path.
|
|
result = self._agent().run_sync(json.dumps(base, ensure_ascii=False))
|
|
return {"ok": True, "model": self.config.model, "outputType": type(result.output).__name__}
|