feat(simulation): add Worker AI polygon runtime and terrain navigation
This commit is contained in:
@@ -0,0 +1,154 @@
|
||||
"""Worker-local, bounded asynchronous inference. No operator-network dependencies."""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class LatestInference:
|
||||
def __init__(self, factory, policy, directory: Path, clock=time.monotonic):
|
||||
self.factory, self.policy, self.directory, self.clock = factory, policy, directory, clock
|
||||
self.condition = threading.Condition()
|
||||
self.enabled = self.closed = self.ready = False
|
||||
self.ever_ready = False
|
||||
self.epoch = self.count = self.dropped = 0
|
||||
self.pending = self.result = None
|
||||
self.error = None
|
||||
self.thread = threading.Thread(target=self._run, name="polygon-inference", daemon=True)
|
||||
|
||||
def start(self):
|
||||
self.thread.start()
|
||||
|
||||
def enable(self, enabled):
|
||||
with self.condition:
|
||||
if enabled != self.enabled:
|
||||
self.enabled = enabled
|
||||
self.epoch += 1
|
||||
self.pending = self.result = None
|
||||
self.condition.notify_all()
|
||||
|
||||
def submit(self, rgb, frame_id, captured_at, simulation_time_ns, observation=None):
|
||||
with self.condition:
|
||||
if not self.enabled or self.closed:
|
||||
return
|
||||
if self.pending is not None:
|
||||
self.dropped += 1
|
||||
self.pending = (rgb, frame_id, captured_at, simulation_time_ns, self.epoch, observation)
|
||||
self.condition.notify_all()
|
||||
|
||||
def command(self, now, frame_deadline=0.5, command_deadline=0.5):
|
||||
with self.condition:
|
||||
result = self.result
|
||||
if not self.enabled:
|
||||
return 0.0, 0.0, "paused", result
|
||||
if self.error:
|
||||
return 0.0, 0.0, "inference-error", result
|
||||
if result is None or now - result["captured_at"] > frame_deadline:
|
||||
return 0.0, 0.0, "stale-camera", result
|
||||
if now - result["completed_at"] > command_deadline:
|
||||
return 0.0, 0.0, "stale-command", result
|
||||
decision = result["decision"]
|
||||
return decision["speed_mps"], decision["yaw_rate_rps"], "none", result
|
||||
|
||||
def _run(self):
|
||||
model = None
|
||||
policy_epoch = None
|
||||
self.directory.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
with (self.directory / "decisions.jsonl").open("a", encoding="utf-8") as journal:
|
||||
while True:
|
||||
with self.condition:
|
||||
self.condition.wait_for(
|
||||
lambda: self.closed or (self.enabled and self.pending is not None)
|
||||
)
|
||||
if self.closed:
|
||||
return
|
||||
rgb, frame_id, captured_at, sim_ns, epoch, observation = self.pending
|
||||
self.pending = None
|
||||
try:
|
||||
if model is None:
|
||||
model = self.factory()
|
||||
model.ready()
|
||||
started = self.clock()
|
||||
evidence = {}
|
||||
if self.policy is None:
|
||||
if policy_epoch is not None and policy_epoch != epoch:
|
||||
model.reset()
|
||||
policy_epoch = epoch
|
||||
decision, boxes, evidence = model.infer_observation(
|
||||
rgb, observation, frame_id
|
||||
)
|
||||
else:
|
||||
road, boxes = model.infer(rgb)
|
||||
with self.condition:
|
||||
if epoch != self.epoch or not self.enabled:
|
||||
continue
|
||||
if self.policy is not None and policy_epoch != epoch:
|
||||
self.policy.reset()
|
||||
policy_epoch = epoch
|
||||
if self.policy is not None:
|
||||
decision = self.policy.decide(road, boxes).model_dump()
|
||||
completed = self.clock()
|
||||
record = dict(
|
||||
frame_id=frame_id,
|
||||
captured_at=captured_at,
|
||||
simulation_time_ns=sim_ns,
|
||||
completed_at=completed,
|
||||
inference_ms=(completed - started) * 1000,
|
||||
decision=decision,
|
||||
)
|
||||
with self.condition:
|
||||
self.ready, self.ever_ready, self.error = True, True, None
|
||||
if epoch != self.epoch or not self.enabled:
|
||||
continue # A pause/resume fences every earlier in-flight decision.
|
||||
self.result = record
|
||||
self.count += 1
|
||||
# Archive source RGB on Worker, independently of the video stream.
|
||||
from PIL import Image
|
||||
|
||||
image_path = self.directory / f"{frame_id:08d}.jpg"
|
||||
Image.fromarray(rgb).save(image_path, "JPEG", quality=90)
|
||||
record = {
|
||||
**record,
|
||||
"image": image_path.name,
|
||||
"image_sha256": hashlib.sha256(image_path.read_bytes()).hexdigest(),
|
||||
"obstacles": boxes,
|
||||
**evidence,
|
||||
}
|
||||
journal.write(json.dumps(record, allow_nan=False) + "\n")
|
||||
journal.flush()
|
||||
except Exception as exc:
|
||||
with (self.directory / "errors.jsonl").open(
|
||||
"a", encoding="utf-8"
|
||||
) as errors:
|
||||
errors.write(
|
||||
json.dumps(
|
||||
{
|
||||
"frame_id": frame_id,
|
||||
"monotonic": self.clock(),
|
||||
"error": str(exc),
|
||||
"type": type(exc).__name__,
|
||||
"traceback": traceback.format_exc(),
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
with self.condition:
|
||||
self.error, self.ready, self.result = type(exc).__name__, False, None
|
||||
if model is not None:
|
||||
model.close()
|
||||
model = None
|
||||
with self.condition:
|
||||
self.condition.wait(timeout=0.5)
|
||||
finally:
|
||||
if model is not None:
|
||||
model.close()
|
||||
|
||||
def close(self):
|
||||
with self.condition:
|
||||
self.closed = True
|
||||
self.condition.notify_all()
|
||||
self.thread.join(timeout=12)
|
||||
Reference in New Issue
Block a user