mirror of
https://github.com/legop3/MultiRoombaRover.git
synced 2026-11-14 10:42:45 -05:00
168 lines
8.2 KiB
Python
168 lines
8.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Live policy worker. JSON state on stdin, proposals on stdout; no rover writes."""
|
|
import argparse
|
|
import collections
|
|
import hashlib
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
import av
|
|
import numpy as np
|
|
import torch
|
|
|
|
from policy import Policy, VERSION, SLOTS, VALUES, HEIGHT, WIDTH, encode_command, sensor_vector
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('--model', required=True)
|
|
parser.add_argument('--url', required=True)
|
|
parser.add_argument('--threads', type=int, default=4)
|
|
parser.add_argument('--threshold', type=float, default=0.7)
|
|
args = parser.parse_args()
|
|
torch.set_num_threads(args.threads)
|
|
torch.set_num_interop_threads(1)
|
|
directory = Path(args.model)
|
|
metadata = json.loads((directory / 'model.json').read_text())
|
|
if metadata['specification']['version'] != VERSION:
|
|
raise ValueError('Model specification is incompatible')
|
|
if hashlib.sha256((directory / 'weights.pt').read_bytes()).hexdigest() != metadata['sha256']:
|
|
raise ValueError('Model weights checksum mismatch')
|
|
schema = metadata['specification']['commands']
|
|
model = Policy().eval()
|
|
model.load_state_dict(torch.load(directory / 'weights.pt', map_location='cpu', weights_only=True))
|
|
observations = collections.deque(maxlen=2048)
|
|
mutex = threading.Lock()
|
|
closed = threading.Event()
|
|
|
|
def read_input():
|
|
try:
|
|
for line in sys.stdin:
|
|
event = json.loads(line)
|
|
with mutex:
|
|
observations.append(event)
|
|
finally:
|
|
closed.set()
|
|
|
|
threading.Thread(target=read_input, daemon=True).start()
|
|
values = np.zeros((SLOTS, VALUES), dtype=np.float32)
|
|
last_commands = np.full(SLOTS, -np.inf)
|
|
sensors, sensor_time = {}, 0
|
|
history = collections.deque(maxlen=32)
|
|
latched = set()
|
|
sampled = 0
|
|
with av.open(args.url, options={'rtsp_transport': 'tcp', 'timeout': '3000000',
|
|
'use_wallclock_as_timestamps': '1'}) as video:
|
|
stream = video.streams.video[0]
|
|
stream.codec_context.thread_count = 1
|
|
for frame in video.decode(stream):
|
|
if closed.is_set():
|
|
return
|
|
if frame.pts is None:
|
|
continue
|
|
stamp = float(frame.pts * frame.time_base) * 1000
|
|
now = time.time() * 1000
|
|
if not math.isfinite(stamp) or abs(now - stamp) > 2000:
|
|
raise ValueError('Video time is stale or not aligned to server time')
|
|
if stamp - sampled < 99:
|
|
continue
|
|
sampled = stamp
|
|
with mutex:
|
|
while observations and observations[0]['ts'] < stamp:
|
|
event = observations.popleft()
|
|
if event['kind'] == 'settings':
|
|
args.threshold = max(.05, min(.99, float(event['threshold'])))
|
|
elif event['kind'] == 'sensor':
|
|
sensors, sensor_time = event['data'], event['ts']
|
|
elif event['kind'] == 'command':
|
|
try:
|
|
encoded = encode_command(event['data'], schema)
|
|
if encoded:
|
|
slot, vector = encoded
|
|
values[slot], last_commands[slot] = vector, event['ts']
|
|
except ValueError:
|
|
pass # Unseen manual controls do not stop autonomy.
|
|
if stamp - sensor_time > 1000:
|
|
history.clear()
|
|
continue
|
|
state = np.concatenate((sensor_vector(sensors), values.flatten(),
|
|
np.minimum((stamp - last_commands) / 10000, 1).astype(np.float32),
|
|
np.isfinite(last_commands).astype(np.float32)))
|
|
pixels = frame.to_ndarray(width=WIDTH, height=HEIGHT, format='rgb24').transpose(2, 0, 1).copy()
|
|
history.append((stamp, pixels, state))
|
|
selected = []
|
|
for offset in (900, 600, 300, 0):
|
|
candidates = [entry for entry in history if entry[0] <= stamp - offset + .01]
|
|
if not candidates or stamp - offset - candidates[-1][0] > 200:
|
|
break
|
|
selected.append(candidates[-1])
|
|
if len(selected) != 4:
|
|
continue
|
|
before = time.perf_counter()
|
|
with torch.inference_mode():
|
|
logits, numeric, wheel_output = model(torch.from_numpy(np.stack([row[1] for row in selected])).unsqueeze(0).float() / 255,
|
|
torch.from_numpy(np.stack([row[2] for row in selected])).unsqueeze(0))
|
|
probabilities = logits[0].sigmoid().tolist()
|
|
predictions = numeric[0].tolist()
|
|
wheel_speeds = [round(value * 500) for value in wheel_output[0].tolist()]
|
|
if not all(math.isfinite(value) for value in probabilities):
|
|
raise ValueError('Nonfinite model output')
|
|
families = {}
|
|
active = set()
|
|
proposals = []
|
|
for slot, descriptor in enumerate(schema):
|
|
if descriptor['template']['type'] == 'drive':
|
|
continue
|
|
command = json.loads(json.dumps(descriptor['template']))
|
|
family = command['type']
|
|
if family == 'peripheral':
|
|
family += ':' + str(command['peripheral']['id']) + ':' + str(command['peripheral']['control'])
|
|
for index, field in enumerate(descriptor['fields']):
|
|
number = predictions[slot][index]
|
|
if not math.isfinite(number):
|
|
raise ValueError('Nonfinite control value')
|
|
number = math.tan(max(-.98, min(.98, number)) * math.pi / 2) * field['scale']
|
|
target = command
|
|
for key in field['path'][:-1]:
|
|
target = target[key]
|
|
target[field['path'][-1]] = round(number) if field['integer'] else number
|
|
proposals.append({'slot': slot, 'score': probabilities[slot], 'command': command,
|
|
'reason': 'Below threshold'})
|
|
if probabilities[slot] < args.threshold:
|
|
continue
|
|
active.add(slot)
|
|
discrete = command['type'] in ('headlight', 'laser', 'horn', 'song', 'raw')
|
|
if family not in families or probabilities[slot] > families[family][0]:
|
|
families[family] = (probabilities[slot], command, slot, discrete)
|
|
latched.intersection_update(active)
|
|
# Choose the winning action before latching, so a held toggle cannot
|
|
# expose a competing lower-confidence toggle on the following tick.
|
|
commands = [command for _, command, slot, discrete in families.values()
|
|
if not discrete or slot not in latched]
|
|
winners = {row[2]: row for row in families.values()}
|
|
for proposal in proposals:
|
|
slot = proposal['slot']
|
|
if slot in active:
|
|
proposal['reason'] = ('Competing action' if slot not in winners else
|
|
'Held discrete action' if winners[slot][3] and slot in latched else 'Proposed')
|
|
for _, _, slot, discrete in families.values():
|
|
if discrete:
|
|
latched.add(slot)
|
|
commands.insert(0, {'type': 'drive', 'driveDirect': dict(zip(('left', 'right'), wheel_speeds))})
|
|
print(json.dumps({'kind': 'prediction', 'frameAt': stamp, 'sensorAt': sensor_time,
|
|
'latencyMs': (time.perf_counter() - before) * 1000,
|
|
'wheelSpeeds': wheel_speeds, 'sensorPresent': sensor_vector(sensors)[32:].tolist(),
|
|
'commands': commands, 'proposals': proposals, 'threshold': args.threshold}, allow_nan=False), flush=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
try:
|
|
main()
|
|
except Exception as error:
|
|
print(json.dumps({'kind': 'error', 'message': str(error)}), flush=True)
|
|
raise SystemExit(1)
|