netris-rl 0.1.0__py3-none-any.whl

netris_rl/__init__.py

@@ -0,0 +1,3 @@
+"""Training utilities for the Netris GameBoy Tetris environment."""
+
+__version__ = "0.1.0"

netris_rl/anchored_ppo.py

@@ -0,0 +1,160 @@
+"""PPO with an auxiliary KL anchor to a frozen reference policy.
+
+Used to fine-tune a behavior-cloned policy without letting RL drift destroy
+it: the total loss is `PPO loss + anchor_coef * KL(anchor || policy)`.
+
+`train()` is adapted from stable-baselines3's PPO (MIT license), with the
+anchor term added.
+"""
+
+import numpy as np
+import torch as th
+import torch.nn.functional as F
+from gymnasium import spaces
+from stable_baselines3 import PPO
+from stable_baselines3.common.utils import explained_variance
+
+
+class AnchoredPPO(PPO):
+    def __init__(self, *args, anchor_model=None, anchor_coef=0.05, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.anchor_coef = anchor_coef
+        self.anchor_policy = None
+        if anchor_model is not None:
+            self.set_anchor(anchor_model)
+
+    def set_anchor(self, anchor_model, coef=None):
+        policy = anchor_model.policy if hasattr(anchor_model, "policy") else anchor_model
+        self.anchor_policy = policy
+        self.anchor_policy.set_training_mode(False)
+        for parameter in self.anchor_policy.parameters():
+            parameter.requires_grad_(False)
+        if coef is not None:
+            self.anchor_coef = coef
+
+    def train(self) -> None:
+        self.policy.set_training_mode(True)
+        self._update_learning_rate(self.policy.optimizer)
+
+        clip_range = self.clip_range(self._current_progress_remaining)
+        clip_range_vf = None
+        if self.clip_range_vf is not None:
+            clip_range_vf = self.clip_range_vf(self._current_progress_remaining)
+
+        entropy_losses = []
+        pg_losses, value_losses = [], []
+        clip_fractions = []
+        anchor_kls = []
+
+        continue_training = True
+        for epoch in range(self.n_epochs):
+            approx_kl_divs = []
+            for rollout_data in self.rollout_buffer.get(self.batch_size):
+                actions = rollout_data.actions
+                if isinstance(self.action_space, spaces.Discrete):
+                    actions = rollout_data.actions.long().flatten()
+
+                values, log_prob, entropy = self.policy.evaluate_actions(
+                    rollout_data.observations, actions
+                )
+                values = values.flatten()
+                advantages = rollout_data.advantages
+                if self.normalize_advantage and len(advantages) > 1:
+                    advantages = (advantages - advantages.mean()) / (
+                        advantages.std() + 1e-8
+                    )
+
+                ratio = th.exp(log_prob - rollout_data.old_log_prob)
+                policy_loss_1 = advantages * ratio
+                policy_loss_2 = advantages * th.clamp(
+                    ratio, 1 - clip_range, 1 + clip_range
+                )
+                policy_loss = -th.min(policy_loss_1, policy_loss_2).mean()
+
+                pg_losses.append(policy_loss.item())
+                clip_fraction = th.mean((th.abs(ratio - 1) > clip_range).float()).item()
+                clip_fractions.append(clip_fraction)
+
+                if clip_range_vf is None:
+                    values_pred = values
+                else:
+                    values_pred = rollout_data.old_values + th.clamp(
+                        values - rollout_data.old_values, -clip_range_vf, clip_range_vf
+                    )
+                value_loss = F.mse_loss(rollout_data.returns, values_pred)
+                value_losses.append(value_loss.item())
+
+                if entropy is None:
+                    entropy_loss = -th.mean(-log_prob)
+                else:
+                    entropy_loss = -th.mean(entropy)
+                entropy_losses.append(entropy_loss.item())
+
+                loss = (
+                    policy_loss
+                    + self.ent_coef * entropy_loss
+                    + self.vf_coef * value_loss
+                )
+
+                if self.anchor_policy is not None:
+                    with th.no_grad():
+                        anchor_dist = self.anchor_policy.get_distribution(
+                            rollout_data.observations
+                        ).distribution
+                    policy_dist = self.policy.get_distribution(
+                        rollout_data.observations
+                    ).distribution
+                    anchor_kl = th.distributions.kl_divergence(
+                        anchor_dist, policy_dist
+                    ).mean()
+                    loss = loss + self.anchor_coef * anchor_kl
+                    anchor_kls.append(float(anchor_kl.item()))
+
+                with th.no_grad():
+                    log_ratio = log_prob - rollout_data.old_log_prob
+                    approx_kl_div = th.mean(
+                        (th.exp(log_ratio) - 1) - log_ratio
+                    ).cpu().numpy()
+                    approx_kl_divs.append(approx_kl_div)
+
+                if self.target_kl is not None and approx_kl_div > 1.5 * self.target_kl:
+                    continue_training = False
+                    if self.verbose >= 1:
+                        print(
+                            f"Early stopping at step {epoch} due to reaching max kl: "
+                            f"{approx_kl_div:.2f}"
+                        )
+                    break
+
+                self.policy.optimizer.zero_grad()
+                loss.backward()
+                th.nn.utils.clip_grad_norm_(
+                    self.policy.parameters(), self.max_grad_norm
+                )
+                self.policy.optimizer.step()
+
+            self._n_updates += 1
+            if not continue_training:
+                break
+
+        explained_var = explained_variance(
+            self.rollout_buffer.values.flatten(),
+            self.rollout_buffer.returns.flatten(),
+        )
+
+        self.logger.record("train/entropy_loss", np.mean(entropy_losses))
+        self.logger.record("train/policy_gradient_loss", np.mean(pg_losses))
+        self.logger.record("train/value_loss", np.mean(value_losses))
+        self.logger.record("train/approx_kl", np.mean(approx_kl_divs))
+        self.logger.record("train/clip_fraction", np.mean(clip_fractions))
+        self.logger.record("train/loss", loss.item())
+        self.logger.record("train/explained_variance", explained_var)
+        if anchor_kls:
+            self.logger.record("train/anchor_kl", np.mean(anchor_kls))
+        if hasattr(self.policy, "log_std"):
+            self.logger.record("train/std", th.exp(self.policy.log_std).mean().item())
+
+        self.logger.record("train/n_updates", self._n_updates, exclude="tensorboard")
+        self.logger.record("train/clip_range", clip_range)
+        if self.clip_range_vf is not None:
+            self.logger.record("train/clip_range_vf", clip_range_vf)

netris_rl/callbacks.py

@@ -0,0 +1,214 @@
+"""Training callbacks: episode metrics, evaluation and best-model selection."""
+
+import csv
+import json
+import math
+import os
+from collections import Counter
+from datetime import UTC, datetime
+from typing import ClassVar
+
+from stable_baselines3.common.callbacks import BaseCallback, EvalCallback
+
+
+class TetrisMetricsCallback(BaseCallback):
+    """Persists per-episode metrics and logs them to TensorBoard.
+
+    Environments put an `episode_summary` dict in `info` when an episode ends
+    (done or truncated). Because envs may live in subprocesses, writing the CSV
+    here (main process) is the only writer.
+    """
+
+    FIELDNAMES: ClassVar[list] = [
+        "timestamp", "episode", "steps", "pieces", "lines_cleared",
+        "max_height", "avg_bumpiness", "avg_holes", "reward_total",
+        "reward_lines", "reward_holes", "reward_height", "reward_bumpiness",
+        "rotations", "avg_placement_col", "avg_placement_height",
+        "action_distribution", "line_clear_distribution",
+    ]
+
+    def __init__(self, log_dir="logs/metrics", verbose=0):
+        super().__init__(verbose)
+        self.log_dir = log_dir
+        os.makedirs(log_dir, exist_ok=True)
+        self.csv_path = os.path.join(log_dir, "episodes.csv")
+        self.episode_buffer = []
+
+    def _on_step(self):
+        infos = self.locals.get("infos")
+        if infos is None:
+            return True
+
+        for info in infos:
+            summary = info.get("episode_summary") if info else None
+            if summary:
+                self.episode_buffer.append(summary)
+                self._record_tensorboard(summary)
+
+        if self.episode_buffer:
+            self._flush_episodes()
+
+        return True
+
+    def _record_tensorboard(self, summary):
+        if self.logger is None:
+            return
+        self.logger.record("train_ep/lines_cleared", summary["lines_cleared"])
+        self.logger.record("train_ep/max_height", summary["max_height"])
+        self.logger.record("train_ep/avg_holes", summary["avg_holes"])
+        self.logger.record("train_ep/avg_bumpiness", summary["avg_bumpiness"])
+        self.logger.record("train_ep/reward_lines", summary["reward_lines"])
+        self.logger.record("train_ep/reward_holes", summary["reward_holes"])
+        self.logger.record("train_ep/reward_height", summary["reward_height"])
+        self.logger.record("train_ep/reward_bumpiness", summary["reward_bumpiness"])
+        try:
+            actions = json.loads(summary.get("action_distribution", "{}"))
+        except (ValueError, TypeError):
+            actions = {}
+        counts = Counter({int(k): int(v) for k, v in actions.items()})
+        total = sum(counts.values())
+        if total:
+            entropy = -sum(
+                (count / total) * math.log(count / total)
+                for count in counts.values() if count
+            )
+            self.logger.record("train_ep/action_entropy", entropy)
+            top_action, top_count = counts.most_common(1)[0]
+            self.logger.record("train_ep/top_action", top_action)
+            self.logger.record("train_ep/top_action_share", top_count / total)
+
+    def _flush_episodes(self):
+        if not self.episode_buffer:
+            return
+        write_header = not os.path.exists(self.csv_path)
+        with open(self.csv_path, "a", newline="") as f:
+            writer = csv.DictWriter(
+                f, fieldnames=self.FIELDNAMES, extrasaction="ignore"
+            )
+            if write_header:
+                writer.writeheader()
+            writer.writerows(self.episode_buffer)
+        self.episode_buffer = []
+
+    def on_training_end(self):
+        self._flush_episodes()
+
+
+class NetrisEvalCallback(EvalCallback):
+    """Evaluation callback that records lines cleared and protects the starter
+    (behavior-cloned) model: a new best is only saved when eval mean lines beat
+    the measured baseline.
+    """
+
+    def __init__(self, *args, run_number=1,
+                 history_path="logs/eval_logs/lines_history.csv",
+                 best_model_dir="logs/best_model", baseline_mean_lines=0.0,
+                 verbose=0, **kwargs):
+        best_model_save_path = kwargs.pop("best_model_save_path", None)
+        super().__init__(
+            *args, verbose=verbose, best_model_save_path=None, **kwargs
+        )
+        if best_model_save_path is not None:
+            best_model_dir = best_model_save_path
+        self.best_model_dir = best_model_dir
+        self.baseline_mean_lines = float(baseline_mean_lines)
+        self.run_number = run_number
+        self.history_path = history_path
+        self.best_eval_mean_lines = self.baseline_mean_lines
+        self.best_eval_max_lines = 0
+        self.last_eval_mean_lines = 0.0
+        self.last_eval_max_lines = 0
+        self.last_eval_mean_reward = 0.0
+        self.last_eval_mean_length = 0.0
+        self._eval_episode_lines = []
+        self._eval_episode_lengths = []
+        os.makedirs(os.path.dirname(history_path), exist_ok=True)
+
+    def _log_success_callback(self, locals_, globals_):
+        super()._log_success_callback(locals_, globals_)
+        info = locals_.get("info") or {}
+        if locals_.get("done"):
+            self._eval_episode_lines.append(int(info.get("lines_cleared", 0)))
+            index = locals_.get("i")
+            lengths = locals_.get("current_lengths")
+            if index is not None and lengths is not None:
+                self._eval_episode_lengths.append(int(lengths[index]))
+
+    def _on_step(self):
+        num_evals_before = len(self.evaluations_timesteps)
+        continue_training = super()._on_step()
+        if len(self.evaluations_timesteps) > num_evals_before:
+            try:
+                self._on_eval_complete()
+            except Exception as e:  # noqa: BLE001
+                print(f"[NetrisEvalCallback] Error reporting eval: {e}")
+        return continue_training
+
+    def _on_eval_complete(self):
+        episode_rewards = self.evaluations_results[-1]
+        mean_reward = float(sum(episode_rewards) / max(1, len(episode_rewards)))
+        lines = self._eval_episode_lines
+        lengths = self._eval_episode_lengths
+        self._eval_episode_lines = []
+        self._eval_episode_lengths = []
+
+        mean_lines = float(sum(lines) / max(1, len(lines)))
+        max_lines = max(lines) if lines else 0
+        mean_length = float(sum(lengths) / max(1, len(lengths))) if lengths else 0.0
+        timestep = int(self.evaluations_timesteps[-1])
+
+        self.last_eval_mean_lines = mean_lines
+        self.last_eval_max_lines = max_lines
+        self.last_eval_mean_reward = mean_reward
+        self.last_eval_mean_length = mean_length
+
+        is_new_best = mean_lines > self.best_eval_mean_lines
+        if is_new_best:
+            self.best_eval_mean_lines = mean_lines
+            self.best_eval_max_lines = max_lines
+            self._save_best_model()
+            print(
+                f"[NetrisEvalCallback] New best: {mean_lines:.2f} mean lines "
+                f"-> {self.best_model_dir}"
+            )
+
+        if self.logger is not None:
+            self.logger.record("eval/mean_lines", mean_lines)
+            self.logger.record("eval/max_lines", max_lines)
+            self.logger.record("eval/mean_length", mean_length)
+            self.logger.record("eval/mean_reward", mean_reward)
+            self.logger.record("eval/best_mean_lines", self.best_eval_mean_lines)
+
+        self._write_history(timestep, lines, mean_reward, mean_length)
+
+    def _save_best_model(self):
+        os.makedirs(self.best_model_dir, exist_ok=True)
+        self.model.save(os.path.join(self.best_model_dir, "best_model"))
+        from stable_baselines3.common.vec_env import VecNormalize
+
+        training_env = self.training_env
+        if isinstance(training_env, VecNormalize):
+            training_env.save(
+                os.path.join(self.best_model_dir, "vecnormalize.pkl")
+            )
+
+    def _write_history(self, timestep, lines, mean_reward, mean_length=0.0):
+        write_header = not os.path.exists(self.history_path)
+        with open(self.history_path, "a", newline="") as f:
+            writer = csv.writer(f)
+            if write_header:
+                writer.writerow([
+                    "timestamp", "timestep", "episodes", "mean_lines",
+                    "max_lines", "mean_reward", "mean_length",
+                    "best_mean_lines",
+                ])
+            writer.writerow([
+                datetime.now(UTC).isoformat(),
+                timestep,
+                len(lines),
+                round(sum(lines) / max(1, len(lines)), 4),
+                max(lines) if lines else 0,
+                round(mean_reward, 4),
+                round(mean_length, 2),
+                round(self.best_eval_mean_lines, 4),
+            ])

netris_rl/cli/bc.py

@@ -0,0 +1,170 @@
+"""Behavior cloning CLI (`netris-bc`).
+
+Trains the CNN policy on teacher/DAgger demonstrations with cross-entropy,
+optionally warm-starting from an existing model.
+"""
+
+import argparse
+import glob
+import os
+import time
+
+import numpy as np
+import torch
+from netris import NetrisEnvConfig, TetrisEnv
+from stable_baselines3 import PPO
+from stable_baselines3.common.vec_env import DummyVecEnv, VecMonitor
+
+from netris_rl.features import TetrisFeaturesExtractor
+
+OBS_KEYS = [
+    "board",
+    "current_piece",
+    "next_piece",
+    "row_completion",
+    "holes",
+    "bumpiness",
+    "aggregate_height",
+]
+
+
+def make_env(rom_path, include_next_piece):
+    return TetrisEnv(
+        rom_path,
+        log_level="error",
+        config=NetrisEnvConfig(include_next_piece=include_next_piece),
+    )
+
+
+def load_demos(patterns, keys=None):
+    keys = keys or OBS_KEYS
+    if isinstance(patterns, str):
+        patterns = [patterns]
+    paths = []
+    for pattern in patterns:
+        paths.extend(glob.glob(pattern))
+    paths = sorted(set(paths))
+    if not paths:
+        raise FileNotFoundError(f"No demo files match {patterns}")
+    parts = [np.load(path) for path in paths]
+    data = {
+        key: np.concatenate([part[key] for part in parts])
+        for key in keys + ["actions"]
+    }
+    print(f"Loaded {len(data['actions'])} demonstrations from {len(paths)} files")
+    return data
+
+
+def batch_observations(data, indices, keys):
+    return {
+        key: torch.as_tensor(data[key][indices], dtype=torch.float32)
+        for key in keys
+    }
+
+
+def accuracy(policy, data, indices, keys, batch_size=1024):
+    correct = 0
+    policy.eval()
+    with torch.no_grad():
+        for start in range(0, len(indices), batch_size):
+            batch = indices[start:start + batch_size]
+            observations = batch_observations(data, batch, keys)
+            distribution = policy.get_distribution(observations)
+            predictions = distribution.distribution.probs.argmax(dim=1)
+            targets = torch.as_tensor(data["actions"][batch], dtype=torch.long)
+            correct += int((predictions == targets).sum())
+    policy.train()
+    return correct / max(1, len(indices))
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Behavior cloning for Netris")
+    parser.add_argument("--demos", nargs="+", default=["logs/demos/demos_*.npz"])
+    parser.add_argument("--rom", default=None,
+                        help="Path to the GameBoy Tetris ROM (or set NETRIS_ROM)")
+    parser.add_argument("--init-model", default="",
+                        help="Warm-start from an existing PPO/BC model")
+    parser.add_argument("--epochs", type=int, default=30)
+    parser.add_argument("--batch-size", type=int, default=256)
+    parser.add_argument("--learning-rate", type=float, default=3e-4)
+    parser.add_argument("--val-fraction", type=float, default=0.1)
+    parser.add_argument("--output", default="logs/bc/bc_model")
+    parser.add_argument("--net-arch", type=int, nargs="+", default=[256, 256])
+    parser.add_argument("--no-next-piece", action="store_true")
+    args = parser.parse_args()
+
+    obs_keys = [
+        key for key in OBS_KEYS
+        if not args.no_next_piece or key != "next_piece"
+    ]
+    data = load_demos(args.demos, obs_keys)
+    total = len(data["actions"])
+    rng = np.random.default_rng(0)
+    permutation = rng.permutation(total)
+    val_size = int(total * args.val_fraction)
+    val_indices = permutation[:val_size]
+    train_indices = permutation[val_size:]
+
+    policy_kwargs = {
+        "features_extractor_class": TetrisFeaturesExtractor,
+        "features_extractor_kwargs": {"features_dim": 256},
+        "net_arch": {"pi": list(args.net_arch), "vf": list(args.net_arch)},
+    }
+
+    env = VecMonitor(DummyVecEnv([
+        lambda: make_env(args.rom, include_next_piece=not args.no_next_piece)
+    ]))
+    if args.init_model:
+        print(f"Warm-starting from {args.init_model}")
+        model = PPO.load(args.init_model, env=env, device="cpu")
+    else:
+        model = PPO(
+            "MultiInputPolicy",
+            env,
+            policy_kwargs=policy_kwargs,
+            device="cpu",
+            verbose=0,
+        )
+    policy = model.policy
+    policy.train()
+    optimizer = torch.optim.Adam(policy.parameters(), lr=args.learning_rate)
+
+    actions = torch.as_tensor(data["actions"], dtype=torch.long)
+    start_time = time.time()
+    val_accuracy = 0.0
+
+    for epoch in range(args.epochs):
+        rng.shuffle(train_indices)
+        epoch_loss = 0.0
+        batches = 0
+        for start in range(0, len(train_indices), args.batch_size):
+            batch = train_indices[start:start + args.batch_size]
+            observations = batch_observations(data, batch, obs_keys)
+            distribution = policy.get_distribution(observations)
+            log_prob = distribution.log_prob(actions[batch])
+            loss = -log_prob.mean()
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()
+            batches += 1
+
+        val_accuracy = accuracy(policy, data, val_indices, obs_keys)
+        print(
+            f"epoch {epoch + 1:02d}/{args.epochs} "
+            f"loss={epoch_loss / max(1, batches):.4f} "
+            f"val_accuracy={val_accuracy:.3f} "
+            f"elapsed={time.time() - start_time:.0f}s",
+            flush=True,
+        )
+
+    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
+    model.save(args.output)
+    print(f"Saved BC model to {args.output}.zip (val accuracy {val_accuracy:.3f})")
+    env.close()
+
+
+if __name__ == "__main__":
+    main()