chenlin
3ad29356c9
集成通用机器人数值接口、本机控制桥、LeRobot 插件和统一键盘遥操作。采用离线 CoACD 全臂碰撞配方 revision 4、局部装配区切分与结构自接触，限制直接关节位姿写入并保留安全看门狗。同步版本号、变更记录、来源许可证和兼容性验证。
464 lines
18 KiB
Python
464 lines
18 KiB
Python
"""Independent loopback broker. No simulation, training or LeRobot imports."""
|
||
|
||
import asyncio
|
||
import contextlib
|
||
import hmac
|
||
import re
|
||
import secrets
|
||
import time
|
||
from collections import deque
|
||
|
||
from aiohttp import WSMsgType, web
|
||
|
||
from . import protocol as p
|
||
from .protocol import RobotError
|
||
|
||
MAX_BYTES = 65536
|
||
PREFIX = "/api/control/v1"
|
||
ORIGIN = re.compile(r"http://(?:127\.0\.0\.1|localhost)(?::[0-9]{1,5})?\Z")
|
||
HOST = re.compile(r"(?:127\.0\.0\.1|localhost)(?::[0-9]{1,5})?\Z")
|
||
BROKER = web.AppKey("broker", object)
|
||
|
||
|
||
class Broker:
|
||
def __init__(self, token):
|
||
if (
|
||
not isinstance(token, str)
|
||
or not 16 <= len(token) <= 4096
|
||
or not all(33 <= ord(c) <= 126 for c in token)
|
||
):
|
||
raise ValueError("控制 token 必须为16–4096字符的可打印 ASCII，且不含空白")
|
||
self.token = token
|
||
self.ws = None
|
||
self.descriptor = None
|
||
self.observation = None
|
||
self.observed_at = 0.0
|
||
self.authorization_generation = 0
|
||
self.blocked_generation = -1
|
||
self.authorized = False
|
||
self.lease = None
|
||
self.last_action_at = 0.0
|
||
self.last_action_seq = 0
|
||
self.rate = deque()
|
||
self.pending = {}
|
||
|
||
def robot(self):
|
||
if self.ws is None or self.ws.closed or self.descriptor is None:
|
||
raise RobotError("DISCONNECTED", "浏览器机器人未连接")
|
||
return self.descriptor
|
||
|
||
def fresh_observation(self):
|
||
self.robot()
|
||
# Paused samples stop advancing by design; do not hide the required play
|
||
# step behind a freshness error merely because the user waited to connect.
|
||
if self.observation is not None and self.observation["paused"]:
|
||
raise RobotError("PAUSED", "仿真已暂停；请先在浏览器点击“播放”，再点击“允许外部控制”")
|
||
if self.observation is None or time.monotonic() - self.observed_at > 0.5:
|
||
raise RobotError(
|
||
"STALE",
|
||
"机器人观测超过 500ms 未更新；请保持仿真页面可见，"
|
||
"检查仿真是否卡顿，再播放并重新允许外部控制",
|
||
)
|
||
return self.observation
|
||
|
||
async def state(self, packet):
|
||
p.record(packet, {"type", "observation", "enabled", "authorizationGeneration"})
|
||
if type(packet["enabled"]) is not bool:
|
||
p.invalid("授权标记必须为 bool")
|
||
generation = p.integer(packet["authorizationGeneration"])
|
||
if generation < self.authorization_generation:
|
||
raise RobotError("STALE", "授权代次倒退")
|
||
observed = p.observation(packet["observation"], self.robot())
|
||
old = self.observation
|
||
if old:
|
||
p.same_identity(observed, old, ("sessionId",))
|
||
if observed["modelEpoch"] < old["modelEpoch"] or (
|
||
observed["modelEpoch"] == old["modelEpoch"]
|
||
and observed["sequence"] < old["sequence"]
|
||
):
|
||
raise RobotError("STALE", "观测代次/序号倒退")
|
||
# Sequence may restart only in a new epoch. Same-sample heartbeats aren't fresh.
|
||
if (
|
||
old is None
|
||
or observed["modelEpoch"] > old["modelEpoch"]
|
||
or observed["sequence"] > old["sequence"]
|
||
):
|
||
self.observation = observed
|
||
self.observed_at = time.monotonic()
|
||
if self.lease and generation != self.authorization_generation:
|
||
# Cancel the previous owner without consuming the new explicit grant.
|
||
await self.revoke("浏览器已重新授权", notify=False)
|
||
self.authorization_generation = generation
|
||
self.authorized = packet["enabled"] and generation > self.blocked_generation
|
||
if self.lease and (
|
||
not self.authorized
|
||
or observed["paused"]
|
||
or any(observed[k] != self.lease[k] for k in ("sessionId", "modelEpoch"))
|
||
):
|
||
await self.revoke("浏览器已撤销控制/重置模型", notify=False)
|
||
|
||
async def rpc(self, operation, payload):
|
||
self.robot()
|
||
if operation == "action":
|
||
for key, (future, op) in list(self.pending.items()):
|
||
if op == "action":
|
||
if not future.done():
|
||
future.set_exception(RobotError("SUPERSEDED", "已由更新的目标替代"))
|
||
self.pending.pop(key, None)
|
||
if len(self.pending) >= 8:
|
||
raise RobotError("CONFLICT", "待确认请求已满")
|
||
request_id = secrets.token_hex(16)
|
||
future = asyncio.get_running_loop().create_future()
|
||
ws, lease, generation = self.ws, self.lease, self.authorization_generation
|
||
self.pending[request_id] = (future, operation)
|
||
try:
|
||
await asyncio.wait_for(
|
||
self.ws.send_json(
|
||
{"type": "request", "id": request_id, "op": operation, "payload": payload}
|
||
),
|
||
1,
|
||
)
|
||
return await asyncio.wait_for(future, 1)
|
||
except TimeoutError as exc:
|
||
if (
|
||
self.ws is ws
|
||
and self.lease == lease
|
||
and self.authorization_generation == generation
|
||
):
|
||
await self.revoke("动作应用确认超时")
|
||
raise RobotError("TIMEOUT", "浏览器未在一秒内确认应用") from exc
|
||
finally:
|
||
self.pending.pop(request_id, None)
|
||
if not future.done():
|
||
future.cancel()
|
||
elif not future.cancelled():
|
||
future.exception() # Also consume failures delivered during a blocked send.
|
||
|
||
async def revoke(self, reason, notify=True):
|
||
self.lease = None
|
||
self.authorized = False
|
||
self.blocked_generation = self.authorization_generation
|
||
for future, _ in self.pending.values():
|
||
if not future.done():
|
||
future.set_exception(RobotError("DISCONNECTED", reason))
|
||
self.pending.clear()
|
||
if notify and self.ws is not None and not self.ws.closed and self.observation:
|
||
with contextlib.suppress(ConnectionError, TimeoutError):
|
||
await asyncio.wait_for(
|
||
self.ws.send_json(
|
||
{
|
||
"type": "stop",
|
||
"reason": reason,
|
||
"sessionId": self.observation["sessionId"],
|
||
"authorizationGeneration": self.authorization_generation,
|
||
}
|
||
),
|
||
0.1,
|
||
)
|
||
|
||
def require_lease(self, request):
|
||
provided = request.headers.get("X-Control-Lease", "")
|
||
if (
|
||
not self.lease
|
||
or not provided.isascii()
|
||
or not hmac.compare_digest(provided.encode(), self.lease["leaseId"].encode())
|
||
):
|
||
raise RobotError("UNAUTHORIZED", "控制租约无效")
|
||
return dict(self.lease)
|
||
|
||
async def monitor(self):
|
||
while True:
|
||
await asyncio.sleep(0.05)
|
||
if self.lease and (
|
||
time.monotonic() - self.last_action_at > 0.5
|
||
or time.monotonic() - self.observed_at > 0.5
|
||
):
|
||
await self.revoke("动作或观测看门狗超时 (500ms)")
|
||
|
||
|
||
@web.middleware
|
||
async def security(request, handler):
|
||
broker = request.app[BROKER]
|
||
try:
|
||
if not HOST.fullmatch(request.headers.get("Host", "")):
|
||
raise RobotError("UNAUTHORIZED", "Host 不被允许")
|
||
origin = request.headers.get("Origin")
|
||
if origin is not None and not ORIGIN.fullmatch(origin):
|
||
raise RobotError("UNAUTHORIZED", "Origin 不被允许")
|
||
if request.query:
|
||
raise RobotError("INVALID_MESSAGE", "禁止 URL 查询参数及 URL 中的 token")
|
||
if request.path != "/ws/control/v1":
|
||
expected = f"Bearer {broker.token}".encode()
|
||
provided = request.headers.get("Authorization", "")
|
||
if not provided.isascii() or not hmac.compare_digest(provided.encode(), expected):
|
||
raise RobotError("UNAUTHORIZED", "需要本机控制 Bearer token")
|
||
elif origin is None:
|
||
raise RobotError("UNAUTHORIZED", "浏览器 WebSocket 必须提供 Origin")
|
||
return await handler(request)
|
||
except RobotError as exc:
|
||
status = {
|
||
"UNAUTHORIZED": 401,
|
||
"DISCONNECTED": 503,
|
||
"TIMEOUT": 504,
|
||
"INVALID_MESSAGE": 400,
|
||
"UNSUPPORTED": 400,
|
||
}.get(exc.code, 409)
|
||
return web.json_response({"error": exc.as_dict()}, status=status)
|
||
except web.HTTPRequestEntityTooLarge:
|
||
return web.json_response(
|
||
{"error": {"code": "INVALID_MESSAGE", "message": "消息不能超过64KiB"}}, status=413
|
||
)
|
||
|
||
|
||
async def body(request):
|
||
if request.content_type != "application/json":
|
||
p.invalid("需要 application/json")
|
||
return p.loads(await request.read())
|
||
|
||
|
||
async def health(request):
|
||
b = request.app[BROKER]
|
||
return web.json_response(
|
||
{
|
||
"protocolVersion": 1,
|
||
"backendConnected": b.descriptor is not None,
|
||
"authorized": b.authorized,
|
||
"hasLease": b.lease is not None,
|
||
}
|
||
)
|
||
|
||
|
||
async def robot(request):
|
||
return web.json_response(request.app[BROKER].robot())
|
||
|
||
|
||
async def observation(request):
|
||
return web.json_response(request.app[BROKER].fresh_observation())
|
||
|
||
|
||
async def claim(request):
|
||
b = request.app[BROKER]
|
||
data = p.record(await body(request), {"sessionId", "modelEpoch", "modelFingerprint"})
|
||
p.text(data["sessionId"])
|
||
p.integer(data["modelEpoch"])
|
||
b.robot()
|
||
observed = b.fresh_observation()
|
||
p.same_identity(data, observed, ("sessionId", "modelEpoch"))
|
||
if data["modelFingerprint"] != b.descriptor["modelFingerprint"]:
|
||
raise RobotError("INCOMPATIBLE_MODEL", "模型指纹已变化")
|
||
if b.lease:
|
||
raise RobotError("CONFLICT", "已有一个 Python 控制者")
|
||
if not b.authorized:
|
||
raise RobotError("UNAUTHORIZED", "请先在浏览器显式允许外部控制")
|
||
lease = {
|
||
"sessionId": data["sessionId"],
|
||
"modelEpoch": data["modelEpoch"],
|
||
"leaseId": secrets.token_hex(24),
|
||
}
|
||
b.lease = lease
|
||
b.last_action_at = time.monotonic()
|
||
b.last_action_seq = 0
|
||
b.rate.clear()
|
||
try:
|
||
result = p.record(
|
||
await b.rpc(
|
||
"claim",
|
||
{
|
||
**lease,
|
||
"modelFingerprint": data["modelFingerprint"],
|
||
"authorizationGeneration": b.authorization_generation,
|
||
},
|
||
),
|
||
p.IDENTITY,
|
||
)
|
||
p.same_identity(result, lease)
|
||
if b.lease != lease:
|
||
raise RobotError("STALE", "租约在申请期间失效")
|
||
return web.json_response(result)
|
||
except BaseException:
|
||
if b.lease == lease:
|
||
await b.revoke("租约申请失败")
|
||
raise
|
||
|
||
|
||
async def action(request):
|
||
b = request.app[BROKER]
|
||
lease = b.require_lease(request)
|
||
b.fresh_observation()
|
||
data = p.action(await body(request), b.robot())
|
||
p.same_identity(data, lease)
|
||
if data["actionSeq"] <= b.last_action_seq:
|
||
raise RobotError("STALE", "拒绝重复/乱序动作")
|
||
now = time.monotonic()
|
||
while b.rate and now - b.rate[0] >= 1:
|
||
b.rate.popleft()
|
||
if len(b.rate) >= 100:
|
||
raise RobotError("CONFLICT", "动作频率不能超过100Hz")
|
||
b.rate.append(now)
|
||
b.last_action_seq = data["actionSeq"]
|
||
b.last_action_at = now
|
||
try:
|
||
result = p.action_result(await b.rpc("action", data), b.robot())
|
||
p.same_identity(result, data, (*p.IDENTITY, "actionSeq"))
|
||
except RobotError as exc:
|
||
if exc.code != "SUPERSEDED" and b.lease == lease:
|
||
await b.revoke("动作确认失败")
|
||
raise
|
||
if b.lease != lease:
|
||
raise RobotError("STALE", "动作确认来自失效租约")
|
||
return web.json_response(result)
|
||
|
||
|
||
async def release(request):
|
||
b = request.app[BROKER]
|
||
lease = b.require_lease(request)
|
||
try:
|
||
await b.rpc("release", lease)
|
||
finally:
|
||
if b.lease == lease:
|
||
await b.revoke("控制者断开连接")
|
||
return web.json_response({"released": True})
|
||
|
||
|
||
async def reset(request):
|
||
b = request.app[BROKER]
|
||
lease = b.require_lease(request)
|
||
p.record(await body(request), set())
|
||
if not b.robot()["capabilities"]["reset"]:
|
||
raise RobotError("UNSUPPORTED", "机器人不支持 reset")
|
||
try:
|
||
result = p.observation(await b.rpc("reset", lease), b.robot())
|
||
if (
|
||
result["sessionId"] != lease["sessionId"]
|
||
or result["modelEpoch"] <= lease["modelEpoch"]
|
||
or not result["paused"]
|
||
):
|
||
p.invalid("reset 未返回新代次的暂停观测")
|
||
finally:
|
||
if b.lease == lease:
|
||
await b.revoke("reset 后需要重新授权")
|
||
return web.json_response(result)
|
||
|
||
|
||
async def websocket(request):
|
||
b = request.app[BROKER]
|
||
ws = web.WebSocketResponse(
|
||
max_msg_size=MAX_BYTES, heartbeat=2, receive_timeout=5, compress=False
|
||
)
|
||
await ws.prepare(request)
|
||
registered = False
|
||
times = deque()
|
||
try:
|
||
# Outer deadline: WS ping/pong must NOT restart the authentication clock.
|
||
first = await asyncio.wait_for(ws.receive(), 5)
|
||
if first.type != WSMsgType.TEXT:
|
||
raise RobotError("UNAUTHORIZED", "必须在5秒内通过认证")
|
||
auth = p.record(p.loads(first.data), {"type", "token", "protocolVersion"})
|
||
p.version(auth["protocolVersion"])
|
||
if (
|
||
auth["type"] != "auth"
|
||
or not isinstance(auth["token"], str)
|
||
or not auth["token"].isascii()
|
||
or not hmac.compare_digest(auth["token"].encode(), b.token.encode())
|
||
):
|
||
raise RobotError("UNAUTHORIZED", "WebSocket token 无效")
|
||
if b.ws is not None:
|
||
raise RobotError("CONFLICT", "已有浏览器后端连接")
|
||
b.ws = ws # Reserve the single backend before any further await.
|
||
await ws.send_json({"type": "authenticated"})
|
||
while True:
|
||
message = (
|
||
await asyncio.wait_for(ws.receive(), 5) if not registered else await ws.receive()
|
||
)
|
||
if message.type != WSMsgType.TEXT:
|
||
break
|
||
now = time.monotonic()
|
||
while times and now - times[0] > 1:
|
||
times.popleft()
|
||
if len(times) >= 256:
|
||
raise RobotError("CONFLICT", "浏览器消息过于频繁")
|
||
times.append(now)
|
||
packet = p.loads(message.data)
|
||
if not isinstance(packet, dict):
|
||
p.invalid("消息必须是对象")
|
||
kind = packet.get("type")
|
||
if kind == "register" and not registered:
|
||
p.record(
|
||
packet,
|
||
{"type", "descriptor", "observation", "enabled", "authorizationGeneration"},
|
||
)
|
||
b.descriptor = p.descriptor(packet["descriptor"])
|
||
await b.state({k: v for k, v in packet.items() if k != "descriptor"})
|
||
registered = True
|
||
await ws.send_json({"type": "ready"})
|
||
elif kind == "state" and registered:
|
||
await b.state(packet)
|
||
elif kind == "result" and registered:
|
||
p.record(packet, {"type", "id", "ok", "value"})
|
||
p.text(packet["id"])
|
||
if type(packet["ok"]) is not bool:
|
||
p.invalid("ok 必须为 bool")
|
||
pending = b.pending.get(packet["id"])
|
||
if pending and not pending[0].done():
|
||
if packet["ok"]:
|
||
pending[0].set_result(packet["value"])
|
||
else:
|
||
error = p.record(packet["value"], {"code", "message"})
|
||
if (
|
||
not isinstance(error["code"], str)
|
||
or not isinstance(error["message"], str)
|
||
or len(error["message"]) > 1024
|
||
):
|
||
p.invalid("错误格式无效")
|
||
pending[0].set_exception(RobotError(error["code"], error["message"]))
|
||
else:
|
||
p.invalid("未注册后端或未知消息类型")
|
||
except (RobotError, TimeoutError, ConnectionError) as exc:
|
||
error = (
|
||
exc
|
||
if isinstance(exc, RobotError)
|
||
else RobotError("DISCONNECTED", "浏览器连接中断/超时")
|
||
)
|
||
if not ws.closed:
|
||
with contextlib.suppress(ConnectionError):
|
||
await ws.send_json({"type": "error", "error": error.as_dict()})
|
||
finally:
|
||
if b.ws is ws:
|
||
await b.revoke("浏览器已断开", notify=False)
|
||
b.ws = b.descriptor = b.observation = None
|
||
b.authorization_generation = 0
|
||
b.blocked_generation = -1
|
||
await ws.close()
|
||
return ws
|
||
|
||
|
||
def create_app(token):
|
||
app = web.Application(middlewares=[security], client_max_size=MAX_BYTES)
|
||
app[BROKER] = Broker(token)
|
||
app.add_routes(
|
||
[
|
||
web.get(f"{PREFIX}/health", health),
|
||
web.get(f"{PREFIX}/robot", robot),
|
||
web.get(f"{PREFIX}/observation", observation),
|
||
web.post(f"{PREFIX}/lease", claim),
|
||
web.delete(f"{PREFIX}/lease", release),
|
||
web.post(f"{PREFIX}/action", action),
|
||
web.post(f"{PREFIX}/reset", reset),
|
||
web.get("/ws/control/v1", websocket),
|
||
]
|
||
)
|
||
|
||
async def lifetime(application):
|
||
b = application[BROKER]
|
||
task = asyncio.create_task(b.monitor())
|
||
yield
|
||
task.cancel()
|
||
with contextlib.suppress(asyncio.CancelledError):
|
||
await task
|
||
await b.revoke("桥接服务关闭")
|
||
if b.ws is not None:
|
||
await b.ws.close()
|
||
|
||
app.cleanup_ctx.append(lifetime)
|
||
return app
|