netris-rl 0.1.0__tar.gz

netris_rl-0.1.0/.gitignore

@@ -0,0 +1,36 @@
+# Python caches and bytecode
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+.hypothesis/
+.ipynb_checkpoints/
+
+# Coverage
+.coverage
+.coverage.*
+htmlcov/
+
+# Virtual environments
+venv/
+.venv/
+env/
+
+# Build artifacts
+dist/
+build/
+*.egg-info/
+
+# Game Boy ROM (copyrighted, never commit)
+*.gb
+*.gb.ram
+
+# Local logs and model artifacts
+logs/
+models/
+*.zip
+*.pkl
+
+# Build
+.pypirc

netris_rl-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Javier Aranda
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

netris_rl-0.1.0/PKG-INFO

@@ -0,0 +1,83 @@
+Metadata-Version: 2.5
+Name: netris-rl
+Version: 0.1.0
+Summary: Behavior cloning, DAgger and anchored PPO training for the netris environment
+Project-URL: Homepage, https://github.com/javi-aranda/netris
+Author: Javier Aranda
+License: MIT License
+        
+        Copyright (c) 2026 Javier Aranda
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Keywords: behavior-cloning,gameboy,ppo,reinforcement-learning,tetris
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Requires-Dist: netris<0.2,>=0.1
+Requires-Dist: stable-baselines3>=2.0
+Requires-Dist: tensorboard>=2.0
+Requires-Dist: torch>=2.0
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: ruff>=0.6; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# netris-rl
+
+Training utilities for the [netris](../netris) Game Boy Tetris environment:
+behavior cloning from a heuristic teacher, DAgger rounds and PPO fine-tuning
+with a KL anchor.
+
+## Install
+
+```bash
+pip install netris-rl
+```
+
+## CLIs
+
+```bash
+# collect teacher demonstrations (optionally DAgger with --policy-model)
+netris-collect --rom /path/to/tetris.gb --steps 30000 --workers 2
+
+# behavior cloning (optionally warm-started with --init-model)
+netris-bc --demos "logs/demos/demos_*.npz" --epochs 30 --output logs/bc/model
+
+# PPO fine-tuning, optionally anchored to a frozen reference policy
+netris-train --rom /path/to/tetris.gb --resume --no-obs-norm \
+    --anchor-model logs/bc/model.zip --total-steps 500000
+
+# evaluation
+netris-eval --model logs/best_model/best_model.zip --episodes 20
+```
+
+## Pipeline used to reach a line-clearing agent
+
+1. `netris-collect` with the analytic Dellacherie teacher.
+2. `netris-bc` (behavior cloning).
+3. `netris-collect --policy-model` (DAgger), then `netris-bc --init-model`.
+4. `netris-train --resume --anchor-model <bc model>` (anchored PPO).
+
+## License
+
+MIT.

netris_rl-0.1.0/README.md

@@ -0,0 +1,39 @@
+# netris-rl
+
+Training utilities for the [netris](../netris) Game Boy Tetris environment:
+behavior cloning from a heuristic teacher, DAgger rounds and PPO fine-tuning
+with a KL anchor.
+
+## Install
+
+```bash
+pip install netris-rl
+```
+
+## CLIs
+
+```bash
+# collect teacher demonstrations (optionally DAgger with --policy-model)
+netris-collect --rom /path/to/tetris.gb --steps 30000 --workers 2
+
+# behavior cloning (optionally warm-started with --init-model)
+netris-bc --demos "logs/demos/demos_*.npz" --epochs 30 --output logs/bc/model
+
+# PPO fine-tuning, optionally anchored to a frozen reference policy
+netris-train --rom /path/to/tetris.gb --resume --no-obs-norm \
+    --anchor-model logs/bc/model.zip --total-steps 500000
+
+# evaluation
+netris-eval --model logs/best_model/best_model.zip --episodes 20
+```
+
+## Pipeline used to reach a line-clearing agent
+
+1. `netris-collect` with the analytic Dellacherie teacher.
+2. `netris-bc` (behavior cloning).
+3. `netris-collect --policy-model` (DAgger), then `netris-bc --init-model`.
+4. `netris-train --resume --anchor-model <bc model>` (anchored PPO).
+
+## License
+
+MIT.

netris_rl-0.1.0/pyproject.toml

@@ -0,0 +1,49 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "netris-rl"
+version = "0.1.0"
+description = "Behavior cloning, DAgger and anchored PPO training for the netris environment"
+readme = "README.md"
+license = { file = "LICENSE" }
+requires-python = ">=3.11"
+authors = [{ name = "Javier Aranda" }]
+keywords = ["reinforcement-learning", "ppo", "behavior-cloning", "tetris", "gameboy"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "netris>=0.1,<0.2",
+    "stable-baselines3>=2.0",
+    "torch>=2.0",
+    "tensorboard>=2.0",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=8.0", "ruff>=0.6"]
+
+[project.urls]
+Homepage = "https://github.com/javi-aranda/netris"
+
+[project.scripts]
+netris-train = "netris_rl.cli.train:main"
+netris-eval = "netris_rl.cli.eval:main"
+netris-collect = "netris_rl.cli.collect:main"
+netris-bc = "netris_rl.cli.bc:main"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/netris_rl"]
+
+[tool.pytest.ini_options]
+markers = ["rom: requires a GameBoy Tetris ROM"]
+addopts = "-m \"not rom\""
+
+[tool.ruff]
+line-length = 100
+target-version = "py311"

netris_rl-0.1.0/src/netris_rl/__init__.py

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

netris_rl-0.1.0/src/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-0.1.0/src/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),
+            ])