mirror of
https://github.com/legop3/MultiRoombaRover.git
synced 2026-11-15 03:02:17 -05:00
354 lines
19 KiB
Python
354 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
"""One finite CPU training job; Node owns scheduling, leases, and shutdown.
|
|
|
|
Usage: python train.py --job /path/to/job.json
|
|
Dependencies are packaged at image build time, never installed at runtime.
|
|
"""
|
|
import argparse
|
|
import copy
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import random
|
|
import resource
|
|
import signal
|
|
import time
|
|
import uuid
|
|
|
|
|
|
def emit(kind, **data):
|
|
print(json.dumps({'kind': kind, **data}, allow_nan=False), flush=True)
|
|
|
|
|
|
def atomic_torch_save(torch, value, path):
|
|
temporary = path.with_suffix('.tmp')
|
|
with temporary.open('wb') as stream:
|
|
torch.save(value, stream)
|
|
stream.flush()
|
|
os.fsync(stream.fileno())
|
|
temporary.replace(path)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('--job')
|
|
parser.add_argument('--check-dependencies', action='store_true', help='Describe the installed production policy and exit without creating storage')
|
|
args = parser.parse_args()
|
|
# --help works even before Python dependencies are installed.
|
|
import numpy as np
|
|
import av
|
|
import torch
|
|
from learning import loss_setup, tensors, losses as policy_losses, evaluate, probe
|
|
from dataset import windows
|
|
from policy import Policy, VERSION, SLOTS, VALUES, specification, is_stop, select_slots, initial_schema
|
|
|
|
if args.check_dependencies:
|
|
emit('dependencies', torch=str(torch.__version__), numpy=str(np.__version__), av=str(av.__version__),
|
|
parameters=sum(p.numel() for p in Policy().parameters()), specification=specification(initial_schema()))
|
|
return
|
|
if not args.job:
|
|
parser.error('--job is required unless --check-dependencies is used')
|
|
|
|
job = json.loads(Path(args.job).read_text())
|
|
root = Path(job['root'])
|
|
config = job['config']
|
|
torch.set_num_threads(config['threads'])
|
|
torch.set_num_interop_threads(1)
|
|
# Training yields CPU priority to the Node/media processes. Thread count is
|
|
# an additional bound, not a claim to reserve specific physical cores.
|
|
os.nice(10)
|
|
stopped = False
|
|
|
|
def stop(_signal, _frame):
|
|
nonlocal stopped
|
|
stopped = True
|
|
|
|
signal.signal(signal.SIGTERM, stop)
|
|
signal.signal(signal.SIGINT, stop)
|
|
started = time.monotonic()
|
|
cancelled = lambda: stopped or time.monotonic() - started >= config['maxJobSeconds']
|
|
model = Policy().cpu()
|
|
optimizer = torch.optim.AdamW(model.parameters(), lr=0.0001)
|
|
# The command-event policy has different inputs, targets and weights. Keep its
|
|
# optimizer lineage separate; reuse recordings with their unchanged partition.
|
|
resume_path = root / 'training' / 'resume-v6.pt'
|
|
schema, steps, examples, last_published = initial_schema(), 0, 0, 0
|
|
trained_rovers = set()
|
|
control_profiles = []
|
|
distinct_windows = {}
|
|
published_windows = 0
|
|
last_probe_steps, learning_probe = 0, None
|
|
if resume_path.exists():
|
|
saved = torch.load(resume_path, map_location='cpu', weights_only=True)
|
|
if saved['version'] != VERSION:
|
|
raise ValueError('Resume schema incompatible; refusing to reinterpret existing weights')
|
|
model.load_state_dict(saved['weights'])
|
|
optimizer.load_state_dict(saved['optimizer'])
|
|
schema = saved['schema']
|
|
steps, examples = saved['steps'], saved['examples']
|
|
last_published = saved['lastPublished']
|
|
trained_rovers = set(saved['trainedRovers'])
|
|
control_profiles = saved.get('controlProfiles', [])
|
|
distinct_windows = saved.get('distinctWindows', {})
|
|
published_windows = saved.get('publishedWindows', 0)
|
|
last_probe_steps = saved.get('lastProbeSteps', 0)
|
|
learning_probe = saved.get('learningProbe')
|
|
cpu_started = time.process_time()
|
|
decode_started = time.monotonic()
|
|
dataset_stats = []
|
|
initial_steps = steps
|
|
pool, validation, accepted, rejected = [], [], [], {}
|
|
observed_windows = {}
|
|
sessions = list(job['sessions'])
|
|
random.shuffle(sessions)
|
|
is_held_out = lambda session_id: int(hashlib.sha256(session_id.encode()).hexdigest()[:8], 16) % 5 == 0
|
|
sessions.sort(key=is_held_out)
|
|
per_session = max(1, config['maxSamples'] // max(1, len(sessions)))
|
|
for session_id in sessions:
|
|
if cancelled():
|
|
break
|
|
candidate_schema = copy.deepcopy(schema)
|
|
samples, seen = [], 0
|
|
stats = {}
|
|
buckets = {key: [] for key in ('stop', 'command', 'none')}
|
|
counts = dict.fromkeys(buckets, 0)
|
|
held_out = is_held_out(session_id)
|
|
try:
|
|
directory = root / 'recordings' / session_id
|
|
for sample in windows(directory, candidate_schema, cancelled, stats, extend=not held_out):
|
|
seen += 1
|
|
# Keep evaluation naturally distributed. Training reserves space
|
|
# for recorded commands and releases so periods without commands
|
|
# cannot consume the entire sample budget. No scripted driving required.
|
|
if held_out:
|
|
reservoir, count, capacity = samples, seen, per_session
|
|
else:
|
|
category = ('stop' if any(sample[2][slot] and is_stop(descriptor)
|
|
for slot, descriptor in enumerate(candidate_schema)) else
|
|
'command' if sample[2].any() else 'none')
|
|
counts[category] += 1
|
|
reservoir, count, capacity = buckets[category], counts[category], max(1, per_session // 3)
|
|
if len(reservoir) < capacity:
|
|
reservoir.append(sample)
|
|
else:
|
|
index = random.randrange(count)
|
|
if index < capacity:
|
|
reservoir[index] = sample
|
|
if not held_out:
|
|
samples = [sample for bucket in buckets.values() for sample in bucket]
|
|
if cancelled():
|
|
break # Do not call a partially decoded recording validated.
|
|
if not samples:
|
|
if held_out and stats.get('unsupportedCommandEvents'):
|
|
dataset_stats.append({**stats, 'partition': 'heldOut', 'sampledWindows': 0,
|
|
'skipReason': 'No supported aligned windows'})
|
|
continue
|
|
raise ValueError('No complete, aligned observation/action windows')
|
|
# Entire sessions stay in one partition across all jobs. Neighboring
|
|
# video windows must never straddle training and validation.
|
|
held_out = is_held_out(session_id)
|
|
if held_out:
|
|
validation.extend(samples)
|
|
else:
|
|
schema = candidate_schema
|
|
pool.extend(samples)
|
|
observed_windows[session_id] = seen
|
|
accepted.append(session_id)
|
|
metadata = json.loads((directory / 'session.json').read_text())
|
|
if not held_out:
|
|
trained_rovers.add(str(metadata['roverId']))
|
|
profile = metadata.get('controlProfile')
|
|
if not held_out and profile is not None and profile not in control_profiles:
|
|
control_profiles.append(profile)
|
|
stats.update({'partition': 'heldOut' if held_out else 'training', 'sampledWindows': len(samples)})
|
|
dataset_stats.append(stats)
|
|
emit('dataset', sessionId=session_id, windows=seen, sampled=len(samples), diagnostics=stats)
|
|
except (ValueError, KeyError, TypeError, OSError, EOFError, IndexError, OverflowError, av.error.FFmpegError) as error:
|
|
rejected[session_id] = str(error)[:500]
|
|
emit('skipped', sessionId=session_id, reason=rejected[session_id])
|
|
if len(pool) < config['minimumSamples'] or not any(sample[2].any() for sample in pool) or cancelled():
|
|
emit('complete', trained=False, accepted=[], rejected=rejected, samples=len(pool),
|
|
diagnostics={'dataset': dataset_stats, 'decodeSeconds': time.monotonic() - decode_started},
|
|
reason='Stopped/time limit' if cancelled() else 'Waiting for enough training-partition samples with human actions')
|
|
return
|
|
|
|
decode_seconds = time.monotonic() - decode_started
|
|
audit_started = time.monotonic()
|
|
numeric_mask, positive_weight = loss_setup(pool, schema)
|
|
training_rows = [row for row in dataset_stats if row.get('partition') == 'training']
|
|
natural_count = sum(row.get('usableWindows', 0) for row in training_rows)
|
|
frequencies = np.sum([row.get('commandEventCounts', [0] * SLOTS) for row in training_rows], axis=0) / max(1, natural_count)
|
|
audit_subset = pool[::max(1, len(pool) // 96)][:96]
|
|
before_training = evaluate(model, audit_subset, schema, job['actionThreshold'], frequencies, config['batchSize'], cancelled)
|
|
before_audit_seconds = time.monotonic() - audit_started
|
|
train_started = time.monotonic()
|
|
gradient_norms = []
|
|
losses = []
|
|
model.train()
|
|
for _ in range(config['passesPerJob']):
|
|
random.shuffle(pool)
|
|
for offset in range(0, len(pool), config['batchSize']):
|
|
if cancelled() or steps - initial_steps >= config['maxStepsPerJob']:
|
|
break
|
|
batch = pool[offset:offset + config['batchSize']]
|
|
images, states, targets, values, available = tensors(batch)
|
|
optimizer.zero_grad(set_to_none=True)
|
|
action_logits, predictions = model(images, states)
|
|
event_loss, value_loss = policy_losses(action_logits, predictions, targets, values, available, numeric_mask, positive_weight)
|
|
loss = event_loss + value_loss
|
|
if not torch.isfinite(loss):
|
|
raise ValueError('Nonfinite training loss; checkpoint not saved')
|
|
loss.backward()
|
|
gradient_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1)
|
|
if not torch.isfinite(gradient_norm):
|
|
raise ValueError('Nonfinite training gradient; checkpoint not saved')
|
|
gradient_norms.append(float(gradient_norm))
|
|
optimizer.step()
|
|
steps += 1
|
|
examples += len(batch)
|
|
losses.append(loss.item())
|
|
if steps % 10 == 0:
|
|
emit('progress', steps=steps, examples=examples, loss=loss.item(),
|
|
losses={'events': event_loss.item(), 'commandValues': value_loss.item()})
|
|
if cancelled() or steps - initial_steps >= config['maxStepsPerJob']:
|
|
break
|
|
if steps == initial_steps:
|
|
emit('complete', trained=False, accepted=[], rejected=rejected, reason='No training steps completed')
|
|
return
|
|
|
|
training_seconds = time.monotonic() - train_started
|
|
evaluation_started = time.monotonic()
|
|
distinct_windows.update(observed_windows)
|
|
unique_count = sum(distinct_windows.values())
|
|
coverage = [row for row in dataset_stats if row.get('partition') == 'heldOut']
|
|
unsupported_events = sum(row.get('unsupportedCommandEvents', 0) for row in coverage)
|
|
unsupported_windows = sum(row.get('unsupportedWindows', 0) for row in coverage)
|
|
evaluation = {'status': 'No supported held-out windows' if unsupported_events and not validation else 'waiting for held-out sessions',
|
|
'windows': len(validation), 'unsupportedCommandEvents': unsupported_events,
|
|
'unsupportedWindows': unsupported_windows}
|
|
if validation and not cancelled():
|
|
measured = evaluate(model, validation, schema, job['actionThreshold'], frequencies, config['batchSize'], cancelled)
|
|
evaluated = measured['windows']
|
|
precision, recall = measured.get('precision', 0), measured.get('recall', 0)
|
|
actual_positive = measured.get('positiveActions', 0)
|
|
mse = measured.get('numericMse', 0)
|
|
stop_actual = measured.get('stopEvents', 0)
|
|
stop_recall = measured.get('stopRecall', 0)
|
|
drive_actual = measured.get('movementEvents', 0)
|
|
drive_recall = measured.get('movementRecall', 0)
|
|
frequency_f1 = max((row['f1'] for row in measured.get('baselines', {}).get('trainingFrequency', [])), default=0)
|
|
held_out_sessions = sum(row.get('partition') == 'heldOut' and not row.get('skipReason') for row in dataset_stats)
|
|
gate_reasons = []
|
|
for condition, reason in [
|
|
(unsupported_events == 0, 'Held-out controls are not fully represented; see unsupportedCommands'),
|
|
(evaluated == len(validation), 'Evaluation interrupted'),
|
|
(held_out_sessions >= 2, 'Need two usable held-out recordings'),
|
|
(evaluated >= 128, 'Need 128 held-out windows'),
|
|
(actual_positive >= 32, 'Need 32 held-out command events'),
|
|
(stop_actual >= 8, 'Need eight held-out stop events'),
|
|
(drive_actual >= 16, 'Need 16 held-out movement commands'),
|
|
(drive_recall >= .5, 'Movement command recall below 0.5'),
|
|
(precision >= .5, 'Command precision below 0.5'),
|
|
(recall >= .5, 'Command recall below 0.5'),
|
|
(stop_recall >= .5, 'Stop recall below 0.5'),
|
|
(measured.get('f1', 0) > frequency_f1, 'Command F1 does not beat the training-frequency baseline'),
|
|
(mse <= .1, 'Command parameter error above 0.1'),
|
|
]:
|
|
if not condition:
|
|
gate_reasons.append(reason)
|
|
evaluation = {**measured, 'status': 'imitation gate not met' if gate_reasons else 'passed imitation gate',
|
|
'unsupportedCommandEvents': unsupported_events, 'unsupportedWindows': unsupported_windows,
|
|
'method': 'Isolated held-out 100ms events; live selection, no temporal latch/physical rollout',
|
|
'threshold': job['actionThreshold'], 'heldOutSessions': held_out_sessions,
|
|
'gateReasons': gate_reasons}
|
|
after_training = evaluate(model, audit_subset, schema, job['actionThreshold'], frequencies, config['batchSize'], cancelled)
|
|
probe_seconds = 0
|
|
if not cancelled() and (learning_probe is None or steps - last_probe_steps >= 1000):
|
|
remaining = config['maxJobSeconds'] - (time.monotonic() - started)
|
|
if remaining > 10:
|
|
probe_started = time.monotonic()
|
|
learning_probe = probe(model, audit_subset, schema, job['actionThreshold'], frequencies,
|
|
cancelled, max_seconds=min(20, remaining / 2))
|
|
learning_probe['atTrainingStep'] = steps
|
|
last_probe_steps = steps
|
|
probe_seconds = time.monotonic() - probe_started
|
|
training_audit = {'before': before_training, 'after': after_training,
|
|
'gradientNormBeforeClipping': {'min': min(gradient_norms), 'max': max(gradient_norms)},
|
|
'positiveWeights': positive_weight[:len(schema)].tolist(),
|
|
'naturalTrainingCommandFrequency': frequencies[:len(schema)].tolist(),
|
|
'sampledCommandFrequency': np.mean([row[2] for row in pool], axis=0)[:len(schema)].tolist(),
|
|
'learningProbe': learning_probe}
|
|
evaluation_seconds = time.monotonic() - evaluation_started - probe_seconds
|
|
publication = {'distinctMinutes': unique_count / 600,
|
|
'newMinutes': (unique_count - published_windows) / 600,
|
|
'minimumMinutes': config['minimumDrivingMinutes'],
|
|
'requiredNewMinutes': config['newDrivingMinutesPerModel']}
|
|
reasons = []
|
|
if publication['distinctMinutes'] < config['minimumDrivingMinutes']:
|
|
reasons.append('Collecting distinct driving footage')
|
|
if publication['newMinutes'] < config['newDrivingMinutesPerModel']:
|
|
reasons.append('Waiting for new driving footage since last publication')
|
|
if steps - last_published < config['checkpointEverySteps']:
|
|
reasons.append('Training steps below publication interval')
|
|
if evaluation['status'] != 'passed imitation gate':
|
|
reasons.append(evaluation['status'])
|
|
if cancelled():
|
|
reasons.append('Job time budget reached')
|
|
publication['reason'] = '; '.join(reasons) or 'Latest updated'
|
|
save_started = time.monotonic()
|
|
published = None
|
|
# Every completed job exports an immutable candidate; Node owns promotion and retention.
|
|
adjectives = ['amber', 'curious', 'gentle', 'quiet', 'silver', 'bright', 'merry', 'sleepy']
|
|
animals = ['otter', 'finch', 'badger', 'robin', 'fox', 'heron', 'marten', 'wren']
|
|
model_id = str(uuid.uuid4())
|
|
name = f'{random.choice(adjectives)}-{random.choice(animals)}-{model_id[:8]}'
|
|
checkpoint_root = root / 'training' / 'checkpoints'
|
|
checkpoint_root.mkdir(exist_ok=True)
|
|
staging = checkpoint_root / f'.{model_id}.pending'
|
|
staging.mkdir()
|
|
weights_path = staging / 'weights.pt'
|
|
atomic_torch_save(torch, model.state_dict(), weights_path)
|
|
published = {
|
|
'id': model_id, 'name': name, 'createdAt': int(time.time() * 1000),
|
|
'trainingSteps': steps, 'examplesProcessed': examples,
|
|
'parameters': sum(p.numel() for p in model.parameters()),
|
|
'weightBytes': weights_path.stat().st_size,
|
|
'sha256': hashlib.sha256(weights_path.read_bytes()).hexdigest(),
|
|
'specification': specification(schema), 'trainedRovers': sorted(trained_rovers),
|
|
'controlProfiles': control_profiles,
|
|
'sourceSessions': accepted, 'trainingLoss': sum(losses) / len(losses),
|
|
'evaluation': evaluation['status'], 'validation': evaluation, 'distinctMinutes': publication['distinctMinutes'], 'dtype': 'float32', 'torchVersion': str(torch.__version__),
|
|
}
|
|
(staging / 'model.json').write_text(json.dumps(published, indent=2, allow_nan=False))
|
|
# A catalog reader sees both weights and metadata, or neither. Existing
|
|
# versions are never overwritten or deleted by the trainer.
|
|
staging.rename(checkpoint_root / model_id)
|
|
if not reasons:
|
|
last_published = steps
|
|
published_windows = unique_count
|
|
atomic_torch_save(torch, {
|
|
'version': VERSION, 'weights': model.state_dict(), 'optimizer': optimizer.state_dict(),
|
|
'schema': schema, 'steps': steps, 'examples': examples,
|
|
'lastPublished': last_published, 'trainedRovers': sorted(trained_rovers),
|
|
'controlProfiles': control_profiles, 'distinctWindows': distinct_windows,
|
|
'publishedWindows': published_windows, 'lastProbeSteps': last_probe_steps, 'learningProbe': learning_probe,
|
|
}, resume_path)
|
|
diagnostics = {'dataset': dataset_stats, 'trainingAudit': training_audit, 'probeSeconds': probe_seconds, 'beforeAuditSeconds': before_audit_seconds, 'decodeSeconds': decode_seconds,
|
|
'trainingSeconds': training_seconds, 'evaluationSeconds': evaluation_seconds, 'checkpointSeconds': time.monotonic() - save_started,
|
|
'cpuSeconds': time.process_time() - cpu_started,
|
|
'peakRssMiB': resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024,
|
|
'dependencies': {'torch': str(torch.__version__), 'numpy': str(np.__version__), 'av': str(av.__version__)},
|
|
'policy': specification(schema), 'browserVideoLatencyMs': None}
|
|
emit('complete', trained=True, accepted=accepted, rejected=rejected,
|
|
steps=steps, examples=examples, loss=sum(losses) / len(losses),
|
|
published=published, promote=not reasons, publication=publication, evaluation=evaluation, diagnostics=diagnostics, elapsedSeconds=time.monotonic() - started)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
try:
|
|
main()
|
|
except Exception as error:
|
|
emit('error', message=str(error))
|
|
raise SystemExit(1)
|