"""Lifecycle owner for the opt-in H2 collector, ordered ball SUB and GPU process.""" from __future__ import annotations from datetime import datetime from dataclasses import replace import multiprocessing as mp from pathlib import Path import queue import signal import time import uuid import numpy as np import torch import yaml from .collector import ResidualCollector from .data import ActionContract, BASE_ACTION_DIM, OBS_DIM, REWARD_PROTOCOL from .events import ADJUDICATION_VERSION, BallStream, CourtGeometry from .journal import AsyncJournal from .learner import learner_main, load_checkpoint from .transport import LearnerStatus, PolicyDoubleBuffer class RLService: def __init__(self, config, h2_contract, spec, origin, *, mode, project_root, checkpoint=None, run_dir=None, observe_only=False, observation_dim=OBS_DIM, policy_metadata=None): if observe_only and (mode != "baseline" or checkpoint is not None): raise ValueError("record-only evaluation requires baseline without a checkpoint") if not isinstance(observation_dim, int) or observation_dim <= 0: raise ValueError("evaluation observation_dim must be a positive integer") if not observe_only and observation_dim != OBS_DIM: raise ValueError("residual service requires the original 182-D observation") if mode not in ("train", "residual", "baseline"): raise ValueError("RL service mode must be train, residual or baseline") if mode == "residual" and not checkpoint: raise ValueError("--rl-mode residual requires --rl-checkpoint") if mode == "baseline" and checkpoint: raise ValueError("baseline mode does not load a checkpoint") inference_policy = None if checkpoint and Path(checkpoint).suffix.lower() == ".onnx": if mode != "residual": raise ValueError("ONNX residual checkpoints are inference-only; use --rl-mode residual") from .onnx_policy import OnnxResidualPolicy inference_policy = OnnxResidualPolicy(checkpoint) saved = inference_policy.saved else: saved = load_checkpoint(checkpoint) if checkpoint else None if saved is not None and mode == "residual": # Evaluate the saved experiment even after the default recipe was # edited. Only transport/output settings belong to this new session. config = replace(type(config)(**saved["config"]), zmq_endpoint=config.zmq_endpoint, output_dir=config.output_dir, episode_queue_capacity=config.episode_queue_capacity, logging_queue_capacity=config.logging_queue_capacity, ball_queue_capacity=config.ball_queue_capacity) self.config, self.mode = config, mode self.observe_only, self.observation_dim = observe_only, observation_dim self.contract = None if observe_only else ActionContract( tuple(h2_contract.policy_names), tuple(h2_contract.action_scale_policy), config.action_scales_rad) if saved is not None: self.contract.validate_saved(saved.get("contract")) if saved is not None and mode == "train" and config.to_dict() != type(config)(**saved["config"]).to_dict(): raise ValueError("training resume requires the checkpoint's exact RL config") parent = Path(config.output_dir) if not parent.is_absolute(): parent = Path(project_root) / parent self.run_dir = Path(run_dir) if run_dir else parent / (datetime.now().strftime("%Y%m%d_%H%M%S") + "_" + uuid.uuid4().hex[:8]) self.run_dir.mkdir(parents=True, exist_ok=False) contract_record = (dict(format="h2-policy-evaluation-v1", robot="h2", obs_dim=observation_dim, action_dim=BASE_ACTION_DIM, policy_names=list(h2_contract.policy_names), action_scale=list(map(float, h2_contract.action_scale_policy)), recording_only=True, action_clip="runner_unchanged") if observe_only else self.contract.to_dict()) resolved = dict(rl=config.to_dict(), contract=contract_record, mode=mode, reward_protocol=REWARD_PROTOCOL, independent_evaluation=mode in ("baseline", "residual"), adjudication_version=ADJUDICATION_VERSION, velocity_source="adapter_zmq", checkpoint=str(Path(checkpoint).resolve()) if checkpoint else None, robot_base_origin_world=list(map(float, origin)), collector_device="cpu", learner_parallel_during_hit=mode == "train", source_commit="4fa48dd2672bd52d9349fb3f7902755d665aa725") if inference_policy is not None: resolved["inference_export"] = saved["export"] if observe_only: resolved.update(recording_only=True, observation_dim=observation_dim, policy=dict(policy_metadata or {}), action_clip="runner_unchanged") (self.run_dir / "resolved_config.yaml").write_text(yaml.safe_dump(resolved, sort_keys=False)) # Avoid tiny CPU MLPs spawning large BLAS/OpenMP teams on each control tick. if not observe_only: torch.set_num_threads(1) self.status = self.sync = self.episodes = self.stop_event = None self.process = self.source = self.journal = self.collector = None self.closed = False self.tick_samples = [] self.source_generation = 0 self.source_error = None self.last_telemetry_s = 0.0 self.control_periods = [] self.last_collector_compute_s = 0.0 self.failure_announced = False previous_signals = {sig: signal.signal(sig, self._interrupt_startup) for sig in (signal.SIGINT, signal.SIGTERM)} try: if mode == "train": context = mp.get_context("spawn") self.status = LearnerStatus(context) self.sync = PolicyDoubleBuffer(context) self.episodes = context.Queue(maxsize=config.episode_queue_capacity) self.stop_event = context.Event() self.journal = AsyncJournal(self.run_dir, config.logging_queue_capacity) self.collector = ResidualCollector(config, self.contract, CourtGeometry.from_spec(spec, origin), self.sync, self.status, self.episodes, self.journal, mode=mode, saved=saved, observe_only=observe_only, observation_dim=observation_dim, inference_policy=inference_policy) if mode == "train": self.process = context.Process(target=learner_main, args=(config, self.contract, self.episodes, self.sync, self.status, self.stop_event, str(self.run_dir), checkpoint), name="h2-residual-learner") started = time.monotonic() print("[STARTUP] learner initialization begin (timeout_s=60)", flush=True) self.process.start() limit = started + 60 while not self.status.ready.value and not self.status.done.value: if self.status.failed.value or not self.process.is_alive(): raise RuntimeError(f"H2 learner startup failed; see {self.run_dir / 'learner_error.txt'}") if time.monotonic() >= limit: raise RuntimeError("H2 learner startup timed out") time.sleep(0.02) print(f"[STARTUP] learner initialization complete elapsed_s={time.monotonic()-started:.3f}", flush=True) self.source = BallStream(config, origin) except BaseException: # No episodes have been collected during startup. A learner stuck # before ready can be stopped without interrupting a training drain. if (self.process is not None and self.process.pid is not None and not self.status.ready.value and not self.status.done.value): self.process.terminate() self.process.join() self.status.failed.value = 1 try: self.close() except Exception as error: print(f"[RL ERROR] startup cleanup: {error}") raise finally: for sig, handler in previous_signals.items(): signal.signal(sig, handler) learner = config.learner_device if mode == "train" else "none" if observe_only: print(f"[POLICY EVAL] run={self.run_dir} obs_dim={observation_dim} recording_only=true action=runner_unchanged") else: print(f"[RL] mode={mode} run={self.run_dir} residual=7x{config.action_scales_rad}rad learner={learner}") if mode != "train": print(f"[RL EVAL] policy={mode} checkpoint={checkpoint or 'none'}; no exploration or learning; stop manually") @staticmethod def _interrupt_startup(signum, _frame): raise KeyboardInterrupt(f"RL startup interrupted by signal {signum}") def _check_process(self): if (self.process is not None and self.process.pid is not None and not self.process.is_alive() and not self.status.done.value): self.status.failed.value = 1 if self.status is not None and self.status.failed.value and not self.failure_announced: print(f"[RL ERROR] learner failed; see {self.run_dir / 'learner_error.txt'}") self.failure_announced = True def poll(self): self.tick_samples = self.source.drain() generation = self.source.error_generation self.source_error = self.source.last_error if generation != self.source_generation else None self.source_generation = generation self._check_process() def step(self, observation, base_action, reference, latest_command_id): now_s = time.monotonic() try: action = self.collector.step(observation, base_action, now_s=now_s, active_command_id=reference.command_id, latest_command_id=latest_command_id, hit_deadline_s=reference.active_hit_deadline_s, target_landing_xy=reference.active_landing_xy, reference_active=not reference.is_frozen, samples=self.tick_samples, racket_distance_m=self.source.racket_distance(now_s), source_fresh=now_s-self.source.last_receipt_s <= self.config.max_ball_gap_s, source_error=self.source_error) except Exception as error: self.collector.disabled = True self.collector.abort(f"rl_runtime_error:{error}") if self.observe_only: print(f"[POLICY EVAL ERROR] scoring disabled, policy unchanged: {error}") else: print(f"[RL ERROR] residual disabled, base policy continues: {error}") action = base_action if self.observe_only else base_action.copy() self.last_collector_compute_s = time.monotonic() - now_s return base_action if self.observe_only else action def applied(self): self.collector.applied() def abort(self, reason): self.collector.abort(reason) def telemetry(self, step_dt, generation_s): now = time.monotonic() self.control_periods.append(float(step_dt)) if now-self.last_telemetry_s >= 1: self.journal.telemetry(dict(monotonic_s=now, control_period_s=float(step_dt), control_period_max_s=max(self.control_periods), control_period_p95_s=float(np.percentile(self.control_periods, 95)), collector_compute_s=self.last_collector_compute_s, closd_generation_s=float(generation_s), **self.collector.snapshot(), journal=self.journal.snapshot(), learner=self.status.snapshot() if self.status else None, learner_backlog_balls=(self.collector.valid_balls-self.status.ingested_balls.value if self.status else 0))) self.last_telemetry_s = now self.control_periods.clear() def close(self): if self.closed: return self.closed = True # The runner has detached from the robot. Further terminal Ctrl+C must # not interrupt checkpoint writes or queue cleanup. previous_signals = {sig: signal.signal(sig, signal.SIG_IGN) for sig in (signal.SIGINT, signal.SIGTERM)} errors = [] try: for resource in (self.collector, self.source): if resource is not None: try: if resource is self.collector: resource.abort("external_stop") else: resource.close() except Exception as error: errors.append(str(error)) if self.process is not None and self.process.pid is not None: self._drain_learner() except Exception as error: errors.append(str(error)) finally: # Failure in one resource must not skip the others. for resource in (self.journal, self.episodes): if resource is not None: try: if resource is self.episodes: if self.status.failed.value: self.episodes.cancel_join_thread() resource.close() if resource is self.episodes and not self.status.failed.value: self.episodes.join_thread() except Exception as error: errors.append(str(error)) for sig, handler in previous_signals.items(): signal.signal(sig, handler) if self.status is not None and self.status.failed.value: errors.append(f"learner failed; recovery artifacts: {self.run_dir}") if errors: raise RuntimeError("RL shutdown incomplete: " + "; ".join(errors)) print(f"[RL EXIT] training and logs complete: {self.run_dir}") def _drain_learner(self): collector = self.collector last_report = -float("inf") stop_sent = False while self.process.is_alive(): self._check_process() collector._try_deliver() if collector.delivery is None and not stop_sent: self.status.stop_after_balls.value = collector.valid_balls self.stop_event.set() stop_sent = True now = time.monotonic() if now - last_report >= 5: backlog = collector.valid_balls - self.status.ingested_balls.value backlog += int(collector.delivery is not None) print(f"[RL DRAIN] balls={backlog} updates={self.status.pending_updates.value} " f"stage={self.status.stage_name} run={self.run_dir}", flush=True) last_report = now self.process.join(timeout=0.1) self.process.join() self._check_process() if self.status.failed.value: self._save_undelivered() def _save_undelivered(self): """Recover queued complete balls after the consumer has exited.""" batches = [] # Queue feeder may still be sending a large episode. A short timed read # lets the feeder advance; unlike join_thread it cannot wait on a dead consumer. while True: try: batches.append(self.episodes.get(timeout=0.2)) except queue.Empty: break delivery = self.collector.delivery if delivery is not None: batches.append(delivery[0]) if batches: path = self.run_dir / "undelivered_episodes.pt" torch.save(batches, path) print(f"[RL RECOVERY] {len(batches)} unconsumed balls saved to {path}") if delivery is not None: _, record = delivery record["summary"].update(delivery_error="learner_failed_before_delivery", recoverable_episode="undelivered_episodes.pt") self.journal.submit_reserved(record) self.collector.delivery = None