Files
MultiRoombaRover/server/src/services/roverLearningService/workers/policy.py
T
2026-09-29 22:14:57 -04:00

200 lines
8.5 KiB
Python

"""Versioned, fixed-size vision/control policy and command encoding.
Slots describe command shapes, not arbitrary executable instructions. Numeric
leaves are regressed; string/bool/list values remain exact template constants.
No command is sent by this module. Execution must use server authorization.
"""
import base64
import itertools
import copy
import json
import math
import numpy as np
import torch
from torch import nn
VERSION = 4
SLOTS = 64
VALUES = 8
FRAMES = 4
HEIGHT, WIDTH = 120, 160
OPERATING_TYPES = {"drive", "motors", "servo", "headlight", "laser", "horn", "song", "raw"}
SENSORS = [
("bumpsAndWheelDrops.bumpLeft", 1), ("bumpsAndWheelDrops.bumpRight", 1),
("bumpsAndWheelDrops.wheelDropLeft", 1), ("bumpsAndWheelDrops.wheelDropRight", 1),
("cliffLeft", 1), ("cliffFrontLeft", 1), ("cliffFrontRight", 1), ("cliffRight", 1),
("wall", 1), ("virtualWall", 1), ("wallSignal", 4096),
("voltageMv", 16000), ("currentMa", 3000), ("batteryChargeMah", 4000),
("batteryCapacityMah", 4000), ("batteryTemperatureC", 50),
("chargingState.code", 5), ("chargingSources.homeBase", 1), ("oiMode.code", 4),
("wheelLeftCurrentMa", 2000), ("wheelRightCurrentMa", 2000),
("mainBrushCurrentMa", 2000), ("sideBrushCurrentMa", 2000),
("lightBumpLeftSignal", 4096), ("lightBumpRightSignal", 4096),
("cliffLeftSignal", 4096), ("cliffRightSignal", 4096),
("distanceMm", 500), ("angleDeg", 180), ("stasis", 1),
]
STATE_SIZE = len(SENSORS) * 2 + SLOTS * (VALUES + 2)
def normalize(value, scale):
return math.atan(float(value) / scale) * 2 / math.pi
def sensor_vector(sensors):
values, present = [], []
for name, scale in SENSORS:
value = sensors
for key in name.split('.'):
value = value.get(key) if isinstance(value, dict) else None
known = isinstance(value, (int, float)) and math.isfinite(value)
values.append(normalize(value, scale) if known else 0)
present.append(float(known))
return np.array(values + present, dtype=np.float32)
def describe_command(command):
if command.get('type') not in OPERATING_TYPES:
return None
if len(json.dumps(command)) > 4096:
raise ValueError('Operating command exceeds 4096-byte schema limit')
values, fields = [], []
def visit(value, path=()):
if isinstance(value, list) and command['type'] == 'horn' and path == ('horn', 'freqs'):
return [visit(item, (*path, index)) for index, item in enumerate(value)]
if isinstance(value, dict):
return {key: visit(item, (*path, key)) for key, item in sorted(value.items())}
# Exact actuator zeros preserve demonstrated releases. Horn frequency
# zeros are parameters so changing a chord does not require a new slot.
if value == 0 and command['type'] != 'horn':
return copy.deepcopy(value)
if isinstance(value, (int, float)) and not isinstance(value, bool):
if not math.isfinite(value) or len(values) >= VALUES:
raise ValueError('Nonfinite command or too many numeric control fields')
scale = {'drive': 500, 'motors': 255, 'servo': 180,
'horn': 1000, 'song': 127}.get(command['type'], 1)
fields.append({'path': list(path), 'scale': scale, 'integer': command['type'] in ('drive', 'motors', 'servo', 'horn') or isinstance(value, int)})
values.append(normalize(value, scale))
return {'$number': len(values) - 1}
# Arrays (e.g. songs) are discrete exact demonstrations, not variable
# length regression outputs. Never reinterpret raw OI bytes as numbers.
return copy.deepcopy(value)
template = visit(command)
return {'template': template, 'fields': fields}, values
def encode_command(command, schema, extend=False):
described = describe_command(command)
if described is None:
return None
descriptor, values = described
try:
slot = schema.index(descriptor)
except ValueError:
if not extend:
raise ValueError('Command shape missing from model schema')
if len(schema) >= SLOTS:
raise ValueError('64 command shapes exhausted; affected windows excluded')
slot = len(schema)
schema.append(descriptor)
vector = np.zeros(VALUES, dtype=np.float32)
vector[:len(values)] = values
return slot, vector
def initial_schema():
# Match commandPipeline.js envelopes, OI_COMMANDS and the four-frequency
# horn control. Values are examples describing fields, not training labels.
commands = []
for left, right in itertools.product((0, 1), repeat=2):
commands.append({'type': 'drive', 'driveDirect': {'left': left, 'right': right}})
for main, side, vacuum in itertools.product((0, 1), repeat=3):
commands.append({'type': 'motors', 'motorPwm': {'main': main, 'side': side, 'vacuum': vacuum}})
for angle in (0, 1):
commands.append({'type': 'servo', 'servo': {'angle': angle}})
for kind in ('headlight', 'laser'):
for action in ('on', 'off', 'toggle'):
commands.append({'type': kind, kind: {'action': action}})
for waveform in ('sine', 'saw'):
commands.append({'type': 'horn', 'horn': {'action': 'start', 'waveform': waveform, 'freqs': [1, 1, 1, 1]}})
commands.append({'type': 'horn', 'horn': {'action': 'stop'}})
for opcode in (128, 131, 132, 143):
commands.append({'type': 'raw', 'raw': base64.b64encode(bytes([opcode])).decode()})
commands.append({'type': 'song', 'song': {'notes': [{'note': 60, 'duration': 8}]}})
schema = []
for command in commands:
encode_command(command, schema, extend=True)
return schema
def command_family(command):
return (command['type'],)
def select_slots(scores, schema, threshold):
winners = {}
for slot, descriptor in enumerate(schema):
if scores[slot] < threshold:
continue
family = command_family(descriptor['template'])
if family not in winners or scores[slot] > scores[winners[family]]:
winners[family] = slot
return set(winners.values())
def decode_command(descriptor, values):
command = copy.deepcopy(descriptor['template'])
for index, field in enumerate(descriptor['fields']):
number = float(values[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
return command
def is_stop(descriptor):
command = descriptor['template']
direct = command.get('driveDirect') or {}
return command['type'] == 'drive' and direct.get('left') == 0 and direct.get('right') == 0
class Policy(nn.Module):
def __init__(self):
super().__init__()
self.vision = nn.Sequential(
nn.Conv2d(3, 24, 5, 2, 2), nn.ReLU(),
nn.Conv2d(24, 48, 3, 2, 1), nn.ReLU(),
nn.Conv2d(48, 96, 3, 2, 1), nn.ReLU(),
nn.Conv2d(96, 128, 3, 2, 1), nn.ReLU(),
nn.AdaptiveAvgPool2d((3, 4)), nn.Flatten(),
nn.Linear(128 * 3 * 4, 128), nn.ReLU(),
)
self.history = nn.GRU(128 + STATE_SIZE, 128, batch_first=True)
self.actions = nn.Linear(128, SLOTS)
self.values = nn.Linear(128, SLOTS * VALUES)
def forward(self, images, state):
batch, frames = images.shape[:2]
encoded = self.vision(images.reshape(-1, 3, HEIGHT, WIDTH)).reshape(batch, frames, -1)
history, _ = self.history(torch.cat((encoded, state), dim=-1))
last = history[:, -1]
return self.actions(last), self.values(last).reshape(batch, SLOTS, VALUES).tanh()
def specification(schema):
return {'version': VERSION, 'architecture': 'spatial-cnn-gru-events-v4',
'frames': FRAMES, 'height': HEIGHT, 'width': WIDTH,
'historyStrideMs': 300, 'decisionIntervalMs': 100,
'stateSize': STATE_SIZE, 'sensorFields': SENSORS,
'normalization': '2/pi * atan(value/scale); missing sensors have mask=0',
'maxCommandShapes': SLOTS, 'numericFieldsPerShape': VALUES,
'commands': schema, 'actionSemantics': 'All controls are command events in the next 100ms; no event preserves existing rover state',
'driveCommandHistoryInput': False, 'zeroNumericValues': 'Exact template constants',
'timing': 'Server-receipt alignment; browser display latency unmeasured'}