| """Independent H2 GPU learner, based on 4fa48dd combined-action residual TD3.""" |
| from __future__ import annotations |
|
|
| import copy |
| import json |
| import os |
| from pathlib import Path |
| import queue |
| import random |
| import signal |
| import time |
| import traceback |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
| from .action import ResidualActionAdapter |
| from .config import RLConfig |
| from .data import ACTOR_INPUT_DIM, OBS_DIM, FORMAT, REWARD_PROTOCOL, ActionContract, EpisodeBatch |
| from .models import Actor, Critic, RunningNormalizer, flatten_parameters |
| from .replay import ReplayBuffer |
|
|
|
|
| def load_checkpoint(path, contract: ActionContract | None = None): |
| |
| saved = torch.load(Path(path), map_location="cpu", weights_only=False) |
| if not isinstance(saved, dict) or saved.get("format") != FORMAT: |
| raise ValueError("requires H2 real residual v3 checkpoint; old H2/G1 checkpoints are unsupported") |
| if saved.get("reward_protocol") != REWARD_PROTOCOL: |
| raise ValueError("checkpoint reward protocol mismatch") |
| if contract is not None: |
| contract.validate_saved(saved.get("contract")) |
| RLConfig(**saved["config"]) |
| return saved |
|
|
|
|
| class TD3Learner: |
| def __init__(self, config: RLConfig, contract: ActionContract, *, device=None): |
| self.config, self.contract = config, contract |
| |
| |
| self.device = torch.device(device or config.learner_device) |
| if self.device.type == "cuda" and not torch.cuda.is_available(): |
| raise RuntimeError("H2 residual learner requires CUDA") |
| random.seed(config.seed) |
| np.random.seed(config.seed) |
| torch.manual_seed(config.seed) |
| self.actor = Actor().to(self.device) |
| self.actor_target = copy.deepcopy(self.actor).eval() |
| self.critics = [Critic().to(self.device), Critic().to(self.device)] |
| self.critic_targets = [copy.deepcopy(c).eval() for c in self.critics] |
| options = dict(betas=config.optimizer_betas, eps=config.optimizer_eps, weight_decay=config.weight_decay) |
| self.actor_optimizer = torch.optim.AdamW(self.actor.parameters(), lr=config.actor_learning_rate, **options) |
| self.critic_optimizers = [torch.optim.AdamW(c.parameters(), lr=config.critic_learning_rate, **options) |
| for c in self.critics] |
| self.normalizer = RunningNormalizer(ACTOR_INPUT_DIM, min_std=config.normalizer_min_std, |
| clip=config.normalizer_clip) |
| self.adapter = ResidualActionAdapter(contract) |
| self.replay = ReplayBuffer(config.replay_capacity, config.seed) |
| self.balls = self.warmup_balls = self.online_balls = self.transitions = 0 |
| self.critic_updates = self.actor_updates = self.warmup_updates = self.online_updates = 0 |
| self.pending_updates = self.policy_version = 0 |
| self.collector_rng_state = None |
| self.last_summary = {} |
| self.metrics = {} |
|
|
| @property |
| def warmup_complete(self): |
| return (self.warmup_balls == self.config.warmup_balls and |
| self.warmup_updates >= self.config.critic_warmup_steps and self.normalizer.frozen) |
|
|
| def ingest(self, episode: EpisodeBatch): |
| episode.validate() |
| if episode.sequence != self.balls + 1: |
| raise ValueError("episode sequence gap; refusing partial/misordered training") |
| expected = "warmup" if self.warmup_balls < self.config.warmup_balls else "online" |
| if episode.phase != expected or (expected == "online" and not self.warmup_complete): |
| raise ValueError("episode training phase mismatch") |
| if self.online_balls >= self.config.online_balls: |
| raise ValueError("training ball budget exceeded") |
| if expected == "warmup": |
| for item in episode.transitions: |
| self.normalizer.update(np.concatenate((item.base_obs, item.base_action))) |
| self.warmup_balls += 1 |
| if self.warmup_balls == self.config.warmup_balls: |
| self.normalizer.freeze() |
| else: |
| self.online_balls += 1 |
| self.pending_updates += len(episode.transitions) * self.config.critic_updates_per_step |
| self.replay.add_episode(episode, self.config.n_step, self.config.discount) |
| self.balls += 1 |
| self.transitions += len(episode.transitions) |
| self.collector_rng_state = episode.collector_rng_state |
| self.last_summary = dict(episode.summary, sent_to_learner=True) |
|
|
| def _soft_update(self, source, target): |
| with torch.no_grad(): |
| for src, dst in zip(source.parameters(), target.parameters()): |
| dst.lerp_(src, self.config.target_tau) |
|
|
| def update_once(self, *, critic_only=False): |
| if self.replay.size < self.config.batch_size: |
| raise RuntimeError("valid warmup balls did not produce a complete replay batch") |
| c = self.config |
| batch = {k: torch.as_tensor(v, device=self.device) for k, v in self.replay.sample(c.batch_size).items()} |
| norm = self.normalizer.normalize_tensor |
| inputs = norm(torch.cat((batch["base_obs"], batch["combined_action"]), dim=1)) |
| with torch.no_grad(): |
| next_inputs = norm(torch.cat((batch["next_base_obs"], batch["next_base_action"]), dim=1)) |
| residual = self.actor_target(next_inputs) |
| if not critic_only: |
| noise = (torch.randn_like(residual) * c.target_noise).clamp(-c.target_noise_clip, c.target_noise_clip) |
| residual = (residual + noise).clamp(-1, 1) |
| combined = self.adapter.combine_tensor(batch["next_base_action"], residual) |
| target_inputs = norm(torch.cat((batch["next_base_obs"], combined), dim=1)) |
| q = torch.minimum(*(target(target_inputs) for target in self.critic_targets)) |
| target = batch["reward_n"] + batch["discount_n"] * batch["bootstrap_mask"] * q |
| losses = [] |
| for critic, optimizer in zip(self.critics, self.critic_optimizers): |
| loss = nn.functional.mse_loss(critic(inputs), target) |
| if not torch.isfinite(loss): |
| raise RuntimeError("nonfinite critic loss") |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| nn.utils.clip_grad_norm_(critic.parameters(), c.gradient_clip, error_if_nonfinite=True) |
| optimizer.step() |
| losses.append(float(loss.detach().cpu())) |
| self.critic_updates += 1 |
| if critic_only: |
| self.warmup_updates += 1 |
| for critic, target_net in zip(self.critics, self.critic_targets): |
| self._soft_update(critic, target_net) |
| else: |
| self.online_updates += 1 |
| self.pending_updates -= 1 |
| if self.online_updates % c.actor_update_delay == 0: |
| for critic in self.critics: |
| critic.requires_grad_(False) |
| actor_inputs = torch.cat((batch["base_obs"], batch["base_action"]), dim=1) |
| residual = self.actor(norm(actor_inputs)) |
| action = self.adapter.combine_tensor(actor_inputs[:, OBS_DIM:], residual) |
| q_loss = -self.critics[0](norm(torch.cat((batch["base_obs"], action), dim=1))).mean() |
| regularization = c.actor_residual_l2 * residual.square().mean() |
| actor_loss = q_loss + regularization |
| if not torch.isfinite(actor_loss): |
| raise RuntimeError("nonfinite actor loss") |
| self.actor_optimizer.zero_grad(set_to_none=True) |
| actor_loss.backward() |
| nn.utils.clip_grad_norm_(self.actor.parameters(), c.gradient_clip, error_if_nonfinite=True) |
| self.actor_optimizer.step() |
| for critic in self.critics: |
| critic.requires_grad_(True) |
| self._soft_update(self.actor, self.actor_target) |
| for critic, target_net in zip(self.critics, self.critic_targets): |
| self._soft_update(critic, target_net) |
| self.actor_updates += 1 |
| self.metrics["actor_loss"] = float(actor_loss.detach().cpu()) |
| self.metrics["actor_q_loss"] = float(q_loss.detach().cpu()) |
| self.metrics["actor_regularization"] = float(regularization.detach().cpu()) |
| self.metrics.update(critic_1_loss=losses[0], critic_2_loss=losses[1]) |
|
|
| def publish(self, sync): |
| self.policy_version = sync.publish(flatten_parameters(self.actor), self.normalizer.mean, |
| self.normalizer.std) |
|
|
| def state_dict(self): |
| counters = {name: getattr(self, name) for name in ( |
| "balls", "warmup_balls", "online_balls", "transitions", "critic_updates", "actor_updates", |
| "warmup_updates", "online_updates", "pending_updates", "policy_version")} |
| return dict(format=FORMAT, reward_protocol=REWARD_PROTOCOL, |
| source_commit="4fa48dd2672bd52d9349fb3f7902755d665aa725", |
| config=self.config.to_dict(), contract=self.contract.to_dict(), counters=counters, |
| actor=self.actor.state_dict(), actor_target=self.actor_target.state_dict(), |
| critics=[c.state_dict() for c in self.critics], |
| critic_targets=[c.state_dict() for c in self.critic_targets], |
| actor_optimizer=self.actor_optimizer.state_dict(), |
| critic_optimizers=[o.state_dict() for o in self.critic_optimizers], |
| normalizer=self.normalizer.state_dict(), replay=self.replay.state_dict(), |
| numpy_rng=np.random.get_state(), python_rng=random.getstate(), torch_rng=torch.get_rng_state(), |
| cuda_rng=torch.cuda.get_rng_state_all() if self.device.type == "cuda" else [], |
| collector_rng=self.collector_rng_state, last_episode=self.last_summary, |
| metrics=self.metrics) |
|
|
| def save(self, path): |
| path = Path(path) |
| temporary = path.with_suffix(path.suffix + ".tmp") |
| with temporary.open("wb") as stream: |
| torch.save(self.state_dict(), stream) |
| stream.flush() |
| os.fsync(stream.fileno()) |
| os.replace(temporary, path) |
|
|
| def restore(self, saved): |
| self.contract.validate_saved(saved["contract"]) |
| if RLConfig(**saved["config"]) != self.config: |
| raise ValueError("resume requires matching resolved H2 training config") |
| self.actor.load_state_dict(saved["actor"]) |
| self.actor_target.load_state_dict(saved["actor_target"]) |
| for group, key in ((self.critics, "critics"), (self.critic_targets, "critic_targets"), |
| (self.critic_optimizers, "critic_optimizers")): |
| for item, state in zip(group, saved[key]): |
| item.load_state_dict(state) |
| self.actor_optimizer.load_state_dict(saved["actor_optimizer"]) |
| norm = saved["normalizer"] |
| self.normalizer.mean = np.asarray(norm["mean"], dtype=np.float64) |
| self.normalizer.m2 = np.asarray(norm["m2"], dtype=np.float64) |
| self.normalizer.count, self.normalizer.frozen = int(norm["count"]), bool(norm["frozen"]) |
| self.replay.load_state_dict(saved["replay"]) |
| for name, value in saved["counters"].items(): |
| setattr(self, name, int(value)) |
| np.random.set_state(saved["numpy_rng"]) |
| random.setstate(saved["python_rng"]) |
| torch.set_rng_state(saved["torch_rng"]) |
| if self.device.type == "cuda" and saved["cuda_rng"]: |
| torch.cuda.set_rng_state_all(saved["cuda_rng"]) |
| self.collector_rng_state = saved["collector_rng"] |
| self.last_summary, self.metrics = saved["last_episode"], saved["metrics"] |
|
|
|
|
| def learner_main(config, contract, episodes, sync, status, stop_event, run_dir, resume_path=None): |
| """Sequential ingestion makes every checkpoint an exact ball/update boundary.""" |
| root = Path(run_dir) |
| learner = None |
| |
| signal.signal(signal.SIGINT, signal.SIG_IGN) |
| try: |
| torch.set_num_threads(1) |
| learner = TD3Learner(config, contract) |
| if resume_path: |
| learner.restore(load_checkpoint(resume_path, contract)) |
| sync.version.value = learner.policy_version |
| learner.publish(sync) |
| last_metrics_s = 0.0 |
|
|
| def progress(): |
| nonlocal last_metrics_s |
| status.ingested_balls.value = learner.balls |
| status.transitions.value = learner.transitions |
| status.critic_updates.value = learner.critic_updates |
| status.actor_updates.value = learner.actor_updates |
| status.pending_updates.value = learner.pending_updates + ( |
| max(0, config.critic_warmup_steps - learner.warmup_updates) |
| if learner.warmup_balls == config.warmup_balls else 0) |
| status.warmup_done.value = int(learner.warmup_complete) |
| now = time.monotonic() |
| if now - last_metrics_s >= 1: |
| with (root / "learner.jsonl").open("a") as log: |
| log.write(json.dumps(dict(monotonic_s=now, **status.snapshot(), **learner.metrics)) + "\n") |
| last_metrics_s = now |
|
|
| |
| progress() |
| status.ready.value = 1 |
| while True: |
| |
| |
| if learner.warmup_balls == config.warmup_balls and not learner.warmup_complete: |
| status.stage.value = 2 |
| while not learner.warmup_complete: |
| started = time.monotonic() |
| learner.update_once(critic_only=True) |
| status.last_update_s.value = time.monotonic() - started |
| progress() |
| learner.publish(sync) |
| if (learner.warmup_complete and learner.online_balls == 0 and |
| not (root / "checkpoint_warmup.pt").exists()): |
| status.stage.value = 4 |
| learner.save(root / "checkpoint_warmup.pt") |
| status.stage.value = 3 |
| while learner.pending_updates: |
| started = time.monotonic() |
| previous_actor_updates = learner.actor_updates |
| learner.update_once() |
| status.last_update_s.value = time.monotonic() - started |
| if learner.actor_updates != previous_actor_updates: |
| learner.publish(sync) |
| progress() |
| progress() |
| if learner.online_balls and learner.online_balls % config.checkpoint_interval_balls == 0: |
| path = root / f"checkpoint_online_ball_{learner.online_balls:06d}.pt" |
| if not path.exists(): |
| status.stage.value = 4 |
| learner.save(path) |
| if learner.online_balls == config.online_balls: |
| break |
| try: |
| status.stage.value = 1 |
| episode = episodes.get(timeout=0.1) |
| except queue.Empty: |
| if stop_event.is_set() and learner.balls >= status.stop_after_balls.value: |
| break |
| continue |
| learner.ingest(episode) |
| progress() |
| status.stage.value = 4 |
| learner.save(root / "checkpoint_final.pt") |
| status.stage.value = 5 |
| status.done.value = 1 |
| except BaseException: |
| status.failed.value = 1 |
| status.stage.value = 6 |
| (root / "learner_error.txt").write_text(traceback.format_exc()) |
| if learner is not None: |
| try: |
| learner.save(root / "checkpoint_interrupted.pt") |
| except Exception: |
| pass |
| finally: |
| status.ready.value = 0 |
|
|
