| """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": |
| |
| |
| 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)) |
| |
| 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: |
| |
| |
| 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 |
| |
| |
| 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: |
| |
| 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 = [] |
| |
| |
| 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 |
|
|
