mirror of
https://github.com/legop3/MultiRoombaRover.git
synced 2026-11-14 18:58:49 -05:00
200 lines
9.7 KiB
Python
200 lines
9.7 KiB
Python
"""Read closed recordings with causal state/action alignment and bounded memory."""
|
|
import bisect
|
|
import collections
|
|
import copy
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import av
|
|
import numpy as np
|
|
|
|
from policy import SLOTS, VALUES, HEIGHT, WIDTH, SENSORS, command_family, is_stop, encode_command, sensor_vector
|
|
|
|
|
|
def load_events(directory):
|
|
metadata = json.loads((directory / 'session.json').read_text())
|
|
if (metadata.get('version') != 1 or not metadata.get('endedAt') or metadata.get('error')
|
|
or metadata.get('droppedEvents') or metadata.get('reason') == 'interrupted'):
|
|
raise ValueError('Incomplete/interrupted recording or lost events')
|
|
if metadata.get('video', {}).get('origin') != 'unix':
|
|
raise ValueError('Unsupported video time origin')
|
|
video_result = metadata.get('video', {}).get('result') or {}
|
|
if video_result.get('signal') == 'SIGKILL' or video_result.get('code') not in (0, 255, None):
|
|
raise ValueError('Video recorder failed or was forcibly killed')
|
|
event_file = directory / 'events.ndjson'
|
|
if event_file.stat().st_size > 64 * 1024 * 1024:
|
|
raise ValueError('Event file exceeds 64 MiB read bound')
|
|
events = []
|
|
with event_file.open() as stream:
|
|
for line in stream:
|
|
event = json.loads(line)
|
|
if event['kind'] == 'captureGap':
|
|
raise ValueError('Recording contains a capture gap')
|
|
if not isinstance(event.get('ts'), (int, float)) or not math.isfinite(event['ts']):
|
|
raise ValueError('Invalid event timestamp')
|
|
events.append(event)
|
|
if not events or not any(e['kind'] == 'sessionEnd' for e in events):
|
|
raise ValueError('Missing completed event stream')
|
|
# File order can differ across rotation buffering. Monotonic time detects
|
|
# wall-clock steps; sorting Unix timestamps alone would hide those steps.
|
|
timed = sorted(events, key=lambda e: int(e['monotonicNs']))
|
|
origin = timed[0]
|
|
for event in timed:
|
|
drift = (event['ts'] - origin['ts']) - (int(event['monotonicNs']) - int(origin['monotonicNs'])) / 1e6
|
|
if abs(drift) > 250:
|
|
raise ValueError('Server wall clock shifted during recording')
|
|
return metadata, sorted(events, key=lambda e: e['ts'])
|
|
|
|
|
|
def prepare(directory, schema, stats):
|
|
metadata, events = load_events(directory)
|
|
requests = {}
|
|
accepted = set()
|
|
failed = set()
|
|
sensors, dispatches, boundaries = [], [], []
|
|
for event in events:
|
|
data = event['data']
|
|
if event['kind'] == 'sensor':
|
|
sensors.append((event['ts'], sensor_vector(data.get('sensors') or {})))
|
|
if event['kind'] in ('driver', 'turn:activeDriver', 'lock', 'private'):
|
|
boundaries.append(event['ts'])
|
|
if event['kind'] != 'command':
|
|
continue
|
|
if data.get('phase') == 'requested' and data.get('source') == 'client':
|
|
requests[data.get('requestId')] = data
|
|
if data.get('phase') == 'result' and data.get('result', {}).get('id'):
|
|
accepted.add(data['result']['id'])
|
|
if data.get('phase') == 'acknowledged' and data.get('acknowledgement', {}).get('error'):
|
|
failed.add(data.get('commandId'))
|
|
if data.get('phase') == 'dispatched':
|
|
dispatches.append(data)
|
|
if not sensors:
|
|
raise ValueError('No sensor observations')
|
|
candidate = copy.deepcopy(schema)
|
|
history, targets = [], []
|
|
stats.update({'sensorFrames': len(sensors), 'dispatches': len(dispatches),
|
|
'failedCommands': len(failed), 'sensorPresentFraction': np.mean([row[1][len(SENSORS):] for row in sensors], axis=0).tolist()})
|
|
for data in dispatches:
|
|
command = data.get('command') or {}
|
|
if command.get('type') in ('reset', 'reboot', 'update', 'sensorStream'):
|
|
boundaries.append(data['ts'])
|
|
encoded = encode_command(command, candidate, extend=True)
|
|
if encoded is None:
|
|
continue
|
|
stamp = data['ts']
|
|
history.append((stamp, *encoded))
|
|
request = requests.get(data.get('requestId'))
|
|
if (data.get('source') == 'client' and request
|
|
and data.get('commandId') in accepted and data.get('commandId') not in failed):
|
|
target_time = request['ts']
|
|
if not 0 <= stamp - target_time <= 1000:
|
|
raise ValueError('Invalid request/dispatch timing')
|
|
targets.append((target_time, *encoded))
|
|
else:
|
|
# Server safety overrides belong in history, never demonstration
|
|
# targets. Exclude windows around them rather than teaching inaction.
|
|
boundaries.append(stamp)
|
|
if not targets:
|
|
raise ValueError('No accepted client operating commands')
|
|
schema[:] = candidate
|
|
return metadata, sensors, sorted(history, key=lambda x: x[0]), sorted(targets, key=lambda x: x[0]), sorted(boundaries)
|
|
|
|
|
|
def windows(directory, schema, cancelled=lambda: False, stats=None):
|
|
directory = Path(directory)
|
|
stats = stats if stats is not None else {}
|
|
metadata, sensors, history, targets, boundaries = prepare(directory, schema, stats)
|
|
stats.update({'usableWindows': 0, 'commandWindows': 0, 'stopWindows': 0, 'noCommandWindows': 0,
|
|
'staleSensorWindows': 0, 'boundaryWindows': 0,
|
|
'ambiguousWindows': 0, 'videoFrames': 0, 'maxVideoGapMs': 0})
|
|
sensor_times = [row[0] for row in sensors]
|
|
target_times = [row[0] for row in targets]
|
|
available = np.zeros(SLOTS, dtype=np.float32)
|
|
for _, slot, _ in history:
|
|
available[slot] = 1
|
|
command_values = np.zeros((SLOTS, VALUES), dtype=np.float32)
|
|
command_times = np.full(SLOTS, -np.inf)
|
|
cursor = 0
|
|
recent = collections.deque(maxlen=32)
|
|
last_sample, last_pts = -math.inf, -math.inf
|
|
def frames():
|
|
files = metadata['video'].get('files', ['video.mkv'])
|
|
for name in files:
|
|
if Path(name).name != name:
|
|
raise ValueError('Invalid video segment path')
|
|
with av.open(str(directory / name)) as container:
|
|
stream = container.streams.video[0]
|
|
stream.codec_context.thread_count = 1
|
|
yield from container.decode(stream)
|
|
|
|
for frame in frames():
|
|
if cancelled():
|
|
return
|
|
if frame.pts is None:
|
|
raise ValueError('Video frame has no timestamp')
|
|
stamp = float(frame.pts * frame.time_base) * 1000
|
|
if not math.isfinite(stamp) or stamp < last_pts:
|
|
raise ValueError('Invalid/nonmonotonic video timestamps')
|
|
stats['videoFrames'] += 1
|
|
if math.isfinite(last_pts):
|
|
stats['maxVideoGapMs'] = max(stats['maxVideoGapMs'], stamp - last_pts)
|
|
last_pts = stamp
|
|
if not metadata['startedAt'] - 2000 <= stamp <= metadata['endedAt'] + 2000:
|
|
raise ValueError('Video PTS is not aligned to session Unix time')
|
|
if stamp < metadata.get('trainingStartedAt', metadata['startedAt']):
|
|
continue
|
|
if stamp - last_sample < 99:
|
|
continue
|
|
last_sample = stamp
|
|
while cursor < len(history) and history[cursor][0] < stamp:
|
|
when, slot, values = history[cursor]
|
|
# Do not let drive predictions simply copy earlier human/model drive commands.
|
|
if schema[slot]['template']['type'] != 'drive':
|
|
command_values[slot] = values
|
|
command_times[slot] = when
|
|
cursor += 1
|
|
sensor_index = bisect.bisect_right(sensor_times, stamp) - 1
|
|
if sensor_index < 0 or stamp - sensor_times[sensor_index] > 1000:
|
|
stats['staleSensorWindows'] += 1
|
|
recent.clear()
|
|
continue
|
|
seen = np.isfinite(command_times).astype(np.float32)
|
|
ages = np.minimum((stamp - command_times) / 10000, 1).astype(np.float32)
|
|
state = np.concatenate((sensors[sensor_index][1], command_values.flatten(), ages, seen))
|
|
pixels = frame.to_ndarray(width=WIDTH, height=HEIGHT, format='rgb24').transpose(2, 0, 1).copy()
|
|
recent.append((stamp, pixels, state))
|
|
chosen = []
|
|
for offset in (900, 600, 300, 0):
|
|
options = [row for row in recent if row[0] <= stamp - offset + 0.01]
|
|
if not options or stamp - offset - options[-1][0] > 200:
|
|
break
|
|
chosen.append(options[-1])
|
|
if len(chosen) != 4 or stamp + 100 > metadata['endedAt']:
|
|
continue
|
|
if bisect.bisect_right(boundaries, stamp + 100) > bisect.bisect_left(boundaries, chosen[0][0]):
|
|
stats['boundaryWindows'] += 1
|
|
continue
|
|
start = bisect.bisect_left(target_times, stamp)
|
|
end = bisect.bisect_left(target_times, stamp + 100)
|
|
actions = np.zeros(SLOTS, dtype=np.float32)
|
|
values = np.zeros((SLOTS, VALUES), dtype=np.float32)
|
|
ambiguous = False
|
|
families = set()
|
|
for _, slot, vector in targets[start:end]:
|
|
template = schema[slot]['template']
|
|
family = command_family(template)
|
|
if ((actions[slot] and (template['type'] in ('headlight', 'laser', 'horn', 'song', 'raw') or not np.array_equal(values[slot], vector)))
|
|
or (not actions[slot] and family in families)):
|
|
ambiguous = True
|
|
break
|
|
families.add(family)
|
|
actions[slot], values[slot] = 1, vector
|
|
if ambiguous:
|
|
stats['ambiguousWindows'] += 1
|
|
continue
|
|
stats['usableWindows'] += 1
|
|
stats['commandWindows' if actions.any() else 'noCommandWindows'] += 1
|
|
stats['stopWindows'] += int(any(actions[slot] and is_stop(descriptor) for slot, descriptor in enumerate(schema)))
|
|
yield np.stack([row[1] for row in chosen]), np.stack([row[2] for row in chosen]), actions, values, available
|