Repository: ai-glimpse/toyrl
Pull request #9: feat: sarsa
Author: login: shenxiangzhuang
Created: 2025-04-27T14:16:49Z | Merged: 2025-04-28T07:14:08Z

Close #7 

<!-- This is an auto-generated comment: release notes by coderabbit.ai -->
## Summary by CodeRabbit

- **New Features**
	- Introduced a SARSA reinforcement learning agent for discrete action environments, featuring neural network-based Q-value estimation and experience replay.
	- Added a configurable training pipeline with support for epsilon-greedy exploration, episodic training, and performance logging.
	- Default support for CartPole-v1 environment, with extensible configuration for other environments.
	- Integrated logging of training metrics such as loss, mean Q-value, and total reward.
- **Documentation**
	- Updated README with new badges, installation instructions, and algorithm implementation checklist.
	- Revised MkDocs configuration with updated repository URL and name.
- **Chores**
	- Adjusted CI workflow to run documentation build only on Ubuntu.
	- Added mypy configuration to skip type checking for wandb imports.
	- Cleaned package exports and removed obsolete hello module.
	- Updated project URLs and simplified dependency groups in configuration.
	- Upgraded pre-commit Ruff linter version.
- **Tests**
	- Added comprehensive test suites for SARSA and Reinforce algorithms covering core components and training loops.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Linked issue #7: SARSA

Diff:
diff --git a/.github/workflows/build_docs.yaml b/.github/workflows/build_docs.yaml
index c9d4c6f..bc0f524 100644
--- a/.github/workflows/build_docs.yaml
+++ b/.github/workflows/build_docs.yaml
@@ -15,8 +15,8 @@ jobs:
           - "3.13"
         os:
           - ubuntu-latest
-          - macos-latest
-          - windows-latest
+          # - macos-latest
+          # - windows-latest
       fail-fast: false
     runs-on: ${{ matrix.os }}
     steps:
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index d62257e..bf5be46 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -32,6 +32,9 @@ jobs:
       - name: Lint with Mypy
         run: uv run mypy .
 
+      # also print ruff version
+      - name: Print Ruff version
+        run: uv run ruff --version
       - name: Lint with Ruff
         run: uv run ruff check
 
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 19238e0..309aec0 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -11,12 +11,14 @@ repos:
         args: [--fix=lf]
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.11.6
+    rev: v0.11.7
     hooks:
       # Run the linter.
       - id: ruff
+        args: ["--config=pyproject.toml"]
       # Run the formatter.
       - id: ruff-format
+        args: ["--config=pyproject.toml"]
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.15.0  # Use the sha / tag you want to point at
     hooks:
diff --git a/README.md b/README.md
index 9c942e9..20cfdab 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,29 @@
 # ToyRL
 
+[![Python](https://img.shields.io/pypi/pyversions/toyrl.svg?color=%2334D058)](https://pypi.org/project/toyrl/)
+[![PyPI](https://img.shields.io/pypi/v/toyrl?color=%2334D058&label=pypi%20package)](https://pypi.org/project/toyrl/)
+[![PyPI Downloads](https://static.pepy.tech/badge/toyrl)](https://pepy.tech/projects/toyrl)
+[![GitHub License](https://img.shields.io/github/license/ai-glimpse/toyrl)](https://github.com/ai-glimpse/toyrl/blob/master/LICENSE)
+
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+[![Checked with mypy](https://www.mypy-lang.org/static/mypy_badge.svg)](https://mypy-lang.org/)
+[![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit)](https://github.com/pre-commit/pre-commit)
+[![Build Docs](https://github.com/ai-glimpse/toyrl/actions/workflows/build_docs.yaml/badge.svg)](https://github.com/ai-glimpse/toyrl/actions/workflows/build_docs.yaml)
+[![Test](https://github.com/ai-glimpse/toyrl/actions/workflows/test.yaml/badge.svg)](https://github.com/ai-glimpse/toyrl/actions/workflows/test.yaml)
+[![Codecov](https://codecov.io/gh/ai-glimpse/toyrl/branch/master/graph/badge.svg)](https://codecov.io/gh/ai-glimpse/toyrl)
+
 ## Documentation
 
-[https://shenxiangzhuang.github.io/toyrl](https://shenxiangzhuang.github.io/toyrl)
+[https://ai-glimpse.github.io/toyrl](https://ai-glimpse.github.io/toyrl)
+
+## Installation
+
+```bash
+pip install toyrl
+```
+
+## Algorithms
+
+- [x] Reinforce
+- [x] SARSA
+- [ ] DQN
diff --git a/mkdocs.yml b/mkdocs.yml
index 7e6b738..ae35670 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -1,6 +1,6 @@
 site_name: ToyRL
-repo_url: https://github.com/shenxiangzhuang/toyrl
-repo_name: shenxiangzhuang/toyrl
+repo_url: https://github.com/ai-glimpse/toyrl
+repo_name: ai-glimpse/toyrl
 site_description: A toy reinforcement learning library
 site_author: Xiangzhuang Shen
 copyright: Copyright &copy; 2025 Xiangzhuang Shen
diff --git a/pyproject.toml b/pyproject.toml
index d2701ee..8f73801 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -34,7 +34,6 @@ dev = [
     "pytest-cov",
     "pytest-sugar",
     "hypothesis>=6.112.0",
-    "commitizen",
 ]
 docs = [
     "mkdocs",
@@ -49,10 +48,10 @@ docs = [
 ]
 
 [project.urls]
-"Homepage" = "https://shenxiangzhuang.github.io/toyrl"
-"Bug Tracker" = "https://shenxiangzhuang.github.io/toyrl/issues"
-"Documentation" = "https://shenxiangzhuang.github.io/toyrl"
-"Source Code" = "https://github.com/shenxiangzhuang/toyrl"
+"Homepage" = "https://ai-glimpse.github.io/toyrl"
+"Bug Tracker" = "https://github.com/ai-glimpse/toyrl/issues"
+"Documentation" = "https://ai-glimpse.github.io/toyrl"
+"Source Code" = "https://github.com/ai-glimpse/toyrl"
 "Release Notes" = "https://github.com/ai-glimpse/toyrl/releases"
 
 [tool.setuptools]
@@ -64,29 +63,20 @@ requires = ["setuptools >= 61.0"]
 build-backend = "setuptools.build_meta"
 
 [dependency-groups]
-dev = [
-    "ipython>=9.1.0",
-]
+dev = ["ipython>=9.1.0"]
 
 
 [tool.setuptools.packages.find]
 include = ["toyrl*"]
 namespaces = false
 
-[tool.commitizen]
-name = "cz_conventional_commits"
-tag_format = "$version"
-version_scheme = "semver"
-version_provider = "poetry"
-update_changelog_on_bump = true
-major_version_zero = true
-changelog_incremental = true
-
-
 [tool.mypy]
 ignore_missing_imports = true
 #strict = true
 
+[[tool.mypy.overrides]]
+module = ["wandb.*"]
+follow_imports = "skip"
 
 # Ruff configuration: https://docs.astral.sh/ruff/configuration/#configuring-ruff
 [tool.ruff]
@@ -135,9 +125,7 @@ select = [
     'D',  # pydocstyle
     'UP', # pyupgrade
 ]
-ignore = [
-    "D"
-]
+ignore = ["D"]
 pydocstyle = { convention = "google" }
 
 # Allow fix for all enabled rules (when `--fix`) is provided.
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..fc6cae6
--- /dev/null
+++ b/tests/__init__.py
@@ -0,0 +1 @@
+"""Test module for toyrl."""
diff --git a/tests/test_hello.py b/tests/test_hello.py
deleted file mode 100644
index 79a7872..0000000
--- a/tests/test_hello.py
+++ /dev/null
@@ -1,6 +0,0 @@
-from toyrl.hello import hello
-
-
-def test_hello() -> None:
-    msg = hello()
-    assert isinstance(msg, str)
diff --git a/tests/test_reinforce.py b/tests/test_reinforce.py
new file mode 100644
index 0000000..47b52a6
--- /dev/null
+++ b/tests/test_reinforce.py
@@ -0,0 +1,141 @@
+"""Tests for the reinforce algorithm."""
+
+import gymnasium as gym
+import numpy as np
+import torch
+
+from toyrl.reinforce import (
+    Agent,
+    EnvConfig,
+    Experience,
+    PolicyNet,
+    ReinforceConfig,
+    ReinforceTrainer,
+    ReplayBuffer,
+    TrainConfig,
+)
+
+
+def test_policy_net():
+    """Test the PolicyNet class."""
+    in_dim, out_dim = 4, 2
+    net = PolicyNet(in_dim=in_dim, out_dim=out_dim)
+
+    # Test forward pass
+    x = torch.randn(in_dim)
+    output = net(x)
+
+    assert output.shape == torch.Size([out_dim])
+    assert isinstance(output, torch.Tensor)
+
+
+def test_replay_buffer():
+    """Test the ReplayBuffer class."""
+    buffer = ReplayBuffer()
+
+    # Test empty buffer
+    assert len(buffer) == 0
+    assert buffer.total_reward() == 0
+
+    # Test adding experience
+    exp = Experience(
+        observation=np.array([0.0, 0.0, 0.0, 0.0]),
+        action=0,
+        action_log_prob=torch.tensor(0.5),
+        reward=1.0,
+        next_observation=np.array([0.1, 0.1, 0.1, 0.1]),
+        terminated=False,
+        truncated=False,
+    )
+    buffer.add_experience(exp)
+
+    assert len(buffer) == 1
+    assert buffer.total_reward() == 1.0
+
+    # Test reset
+    buffer.reset()
+    assert len(buffer) == 0
+
+
+def test_agent_creation():
+    """Test creating an agent."""
+    in_dim, out_dim = 4, 2
+    policy_net = PolicyNet(in_dim=in_dim, out_dim=out_dim)
+    optimizer = torch.optim.Adam(policy_net.parameters(), lr=0.01)
+
+    agent = Agent(policy_net=policy_net, optimizer=optimizer)
+
+    # Check agent properties
+    assert hasattr(agent, "_policy_net")
+    assert hasattr(agent, "_optimizer")
+    assert hasattr(agent, "_replay_buffer")
+
+
+def test_agent_act():
+    """Test the agent's act method."""
+    in_dim, out_dim = 4, 2
+    policy_net = PolicyNet(in_dim=in_dim, out_dim=out_dim)
+    optimizer = torch.optim.Adam(policy_net.parameters(), lr=0.01)
+
+    agent = Agent(policy_net=policy_net, optimizer=optimizer)
+
+    # Test act method
+    observation = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32)
+    action, action_log_prob = agent.act(observation)
+
+    assert isinstance(action, int)
+    assert action in [0, 1]  # For CartPole
+    assert isinstance(action_log_prob, torch.Tensor)
+
+
+def test_config():
+    """Test the Config class."""
+    config = ReinforceConfig()
+
+    # Check default values
+    assert config.env.env_name == "CartPole-v1"
+    assert config.train.gamma == 0.999
+
+    # Test custom config
+    custom_config = ReinforceConfig(
+        env=EnvConfig(env_name="MountainCar-v0", solved_threshold=90.0),
+        train=TrainConfig(num_episodes=1000, learning_rate=0.005),
+    )
+
+    assert custom_config.env.env_name == "MountainCar-v0"
+    assert custom_config.env.solved_threshold == 90.0
+    assert custom_config.train.num_episodes == 1000
+    assert custom_config.train.learning_rate == 0.005
+
+
+def test_trainer_creation():
+    """Test creating a trainer."""
+    config = ReinforceConfig(
+        env=EnvConfig(env_name="CartPole-v1", render_mode=None),
+        train=TrainConfig(num_episodes=10, learning_rate=0.01, log_wandb=False),
+    )
+
+    trainer = ReinforceTrainer(config)
+
+    assert isinstance(trainer.env, gym.Env)
+    assert isinstance(trainer.agent, Agent)
+    assert hasattr(trainer, "num_episodes")
+    assert trainer.num_episodes == 10
+
+
+def test_minimal_training():
+    """Test minimal training run with a single episode."""
+    # Create minimal config with just one episode
+    config = ReinforceConfig(
+        env=EnvConfig(env_name="CartPole-v1", render_mode=None),
+        train=TrainConfig(num_episodes=1, learning_rate=0.01, log_wandb=False),
+    )
+
+    # Initialize trainer
+    trainer = ReinforceTrainer(config)
+
+    # Run training for one episode
+    trainer.train()
+
+    # If we got here without errors, test passed
+    assert True
diff --git a/tests/test_sarsa.py b/tests/test_sarsa.py
new file mode 100644
index 0000000..9fcb57a
--- /dev/null
+++ b/tests/test_sarsa.py
@@ -0,0 +1,164 @@
+"""Tests for the SARSA algorithm."""
+
+import gymnasium as gym
+import numpy as np
+import torch
+
+from toyrl.sarsa import (
+    Agent,
+    EnvConfig,
+    Experience,
+    PolicyNet,
+    ReplayBuffer,
+    SarsaConfig,
+    SarsaTrainer,
+    TrainConfig,
+)
+
+
+def test_policy_net():
+    """Test the PolicyNet class."""
+    env_dim, action_dim, action_num = 4, 1, 2
+    net = PolicyNet(env_dim=env_dim, action_dim=action_dim, action_num=action_num)
+
+    # Test forward pass
+    x = torch.cat((torch.randn(env_dim), torch.tensor([0], dtype=torch.float32)))
+    output = net(x)
+
+    assert output.shape == torch.Size([1])
+    assert isinstance(output, torch.Tensor)
+
+
+def test_replay_buffer():
+    """Test the ReplayBuffer class."""
+    buffer = ReplayBuffer()
+
+    # Test empty buffer
+    assert len(buffer) == 0
+    assert buffer.total_reward() == 0
+
+    # Test adding experience
+    exp = Experience(
+        observation=np.array([0.0, 0.0, 0.0, 0.0]),
+        action=0,
+        reward=1.0,
+        next_observation=np.array([0.1, 0.1, 0.1, 0.1]),
+        next_action=1,
+        terminated=False,
+        truncated=False,
+    )
+    buffer.add_experience(exp)
+
+    assert len(buffer) == 1
+    assert buffer.total_reward() == 1.0
+
+    # Add a second experience
+    exp2 = Experience(
+        observation=np.array([0.1, 0.1, 0.1, 0.1]),
+        action=1,
+        reward=0.5,
+        next_observation=np.array([0.2, 0.2, 0.2, 0.2]),
+        terminated=False,
+        truncated=False,
+    )
+    buffer.add_experience(exp2)
+
+    # Test sampling with next state-action pairs
+    samples = buffer.sample(with_next_sa=True)
+    assert len(samples) == 1
+    assert samples[0].action == 0
+    assert samples[0].next_action == 1
+
+    # Test reset
+    buffer.reset()
+    assert len(buffer) == 0
+
+
+def test_agent_creation():
+    """Test creating an agent."""
+    env_dim, action_dim, action_num = 4, 1, 2
+    policy_net = PolicyNet(env_dim=env_dim, action_dim=action_dim, action_num=action_num)
+    optimizer = torch.optim.RMSprop(policy_net.parameters(), lr=0.01)
+
+    agent = Agent(policy_net=policy_net, optimizer=optimizer)
+
+    # Check agent properties
+    assert hasattr(agent, "_policy_net")
+    assert hasattr(agent, "_optimizer")
+    assert hasattr(agent, "_replay_buffer")
+    assert agent._action_num == action_num
+
+
+def test_agent_act():
+    """Test the agent's act method."""
+    env_dim, action_dim, action_num = 4, 1, 2
+    policy_net = PolicyNet(env_dim=env_dim, action_dim=action_dim, action_num=action_num)
+    optimizer = torch.optim.RMSprop(policy_net.parameters(), lr=0.01)
+
+    agent = Agent(policy_net=policy_net, optimizer=optimizer)
+
+    # Test act method with epsilon=0 (greedy)
+    observation = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32)
+    action, q_value = agent.act(observation, epsilon=0.0)
+
+    assert isinstance(action, int)
+    assert action in [0, 1]  # For CartPole
+    assert isinstance(q_value, float) or q_value is None
+
+    # Test act method with epsilon=1.0 (random)
+    action, q_value = agent.act(observation, epsilon=1.0)
+    assert action in [0, 1]
+    assert q_value is None
+
+
+def test_config():
+    """Test the Config class."""
+    config = SarsaConfig()
+
+    # Check default values
+    assert config.env.env_name == "CartPole-v1"
+    assert config.train.gamma == 0.999
+
+    # Test custom config
+    custom_config = SarsaConfig(
+        env=EnvConfig(env_name="MountainCar-v0", solved_threshold=90.0),
+        train=TrainConfig(num_episodes=1000, learning_rate=0.005),
+    )
+
+    assert custom_config.env.env_name == "MountainCar-v0"
+    assert custom_config.env.solved_threshold == 90.0
+    assert custom_config.train.num_episodes == 1000
+    assert custom_config.train.learning_rate == 0.005
+
+
+def test_trainer_creation():
+    """Test creating a trainer."""
+    config = SarsaConfig(
+        env=EnvConfig(env_name="CartPole-v1", render_mode=None),
+        train=TrainConfig(num_episodes=10, learning_rate=0.01, log_wandb=False),
+    )
+
+    trainer = SarsaTrainer(config)
+
+    assert isinstance(trainer.env, gym.Env)
+    assert isinstance(trainer.agent, Agent)
+    assert hasattr(trainer, "num_episodes")
+    assert trainer.num_episodes == 10
+
+
+def test_minimal_training():
+    """Test minimal training run with a single episode."""
+    # Create minimal config with just one episode
+    config = SarsaConfig(
+        env=EnvConfig(env_name="CartPole-v1", render_mode=None),
+        train=TrainConfig(num_episodes=1, learning_rate=0.01, log_wandb=False),
+    )
+
+    # Initialize trainer
+    trainer = SarsaTrainer(config)
+
+    # Run training for one episode
+    trainer.train()
+
+    # If we got here without errors, test passed
+    assert True
diff --git a/toyrl/__init__.py b/toyrl/__init__.py
index 21a4035..2e3f651 100644
--- a/toyrl/__init__.py
+++ b/toyrl/__init__.py
@@ -1 +1,4 @@
-"""MPPT Module."""
+from .reinforce import ReinforceConfig, ReinforceTrainer
+from .sarsa import SarsaConfig, SarsaTrainer
+
+__all__ = ["ReinforceTrainer", "ReinforceConfig", "SarsaTrainer", "SarsaConfig"]
diff --git a/toyrl/hello.py b/toyrl/hello.py
deleted file mode 100644
index 536b28e..0000000
--- a/toyrl/hello.py
+++ /dev/null
@@ -1,33 +0,0 @@
-"""Hello World module."""
-
-
-def hello() -> str:
-    """Returns a greeting.
-
-    Returns:
-        str: Hello, World!
-
-    Examples:
-        >>> hello()
-        'Hello, World!'
-    """
-    return "Hello, World!"
-
-
-def add(a: float, b: float) -> float:
-    """Adds two numbers together.
-
-    Args:
-        a: First number
-        b: Second number
-
-    Returns:
-        The sum of a and b
-
-    Examples:
-        >>> add(1, 2)
-        3
-        >>> add(2, 3)
-        5
-    """
-    return a + b
diff --git a/toyrl/reinforce.py b/toyrl/reinforce.py
index d7c7f56..f195a5c 100644
--- a/toyrl/reinforce.py
+++ b/toyrl/reinforce.py
@@ -117,16 +117,18 @@ class TrainConfig:
     num_episodes: int = 500
     learning_rate: float = 0.01
     with_baseline: bool = True
+    log_wandb: bool = False
 
 
 @dataclass
-class Config:
+class ReinforceConfig:
     env: EnvConfig = field(default_factory=EnvConfig)
     train: TrainConfig = field(default_factory=TrainConfig)
 
 
 class ReinforceTrainer:
-    def __init__(self, config: Config) -> None:
+    def __init__(self, config: ReinforceConfig) -> None:
+        self.config = config
         self.env = gym.make(config.env.env_name, render_mode=config.env.render_mode)
         env_in_dim = self.env.observation_space.shape[0]  # type: ignore[index]
         env_out_dim = self.env.action_space.n  # type: ignore[attr-defined]
@@ -139,14 +141,14 @@ def __init__(self, config: Config) -> None:
         self.gamma = config.train.gamma
         self.with_baseline = config.train.with_baseline
         self.solved_threshold = config.env.solved_threshold
-
-        wandb.init(
-            # set the wandb project where this run will be logged
-            project="Reinforce",
-            name=f"[{config.env.env_name}]lr={config.train.learning_rate}, baseline={config.train.with_baseline}",
-            # track hyperparameters and run metadata
-            config=asdict(config),
-        )
+        if config.train.log_wandb:
+            wandb.init(
+                # set the wandb project where this run will be logged
+                project="Reinforce",
+                name=f"[{config.env.env_name}]lr={config.train.learning_rate}, baseline={config.train.with_baseline}",
+                # track hyperparameters and run metadata
+                config=asdict(config),
+            )
 
     def train(self) -> None:
         for epi in range(self.num_episodes):
@@ -175,19 +177,20 @@ def train(self) -> None:
             solved = total_reward > self.solved_threshold
             self.agent.onpolicy_reset()
             print(f"Episode {epi}, loss: {loss}, total_reward: {total_reward}, solved: {solved}")
-            wandb.log(
-                {
-                    "episode": epi,
-                    "loss": loss,
-                    "total_reward": total_reward,
-                }
-            )
+            if self.config.train.log_wandb:
+                wandb.log(
+                    {
+                        "episode": epi,
+                        "loss": loss,
+                        "total_reward": total_reward,
+                    }
+                )
 
 
 if __name__ == "__main__":
-    default_config = Config(
+    default_config = ReinforceConfig(
         env=EnvConfig(env_name="CartPole-v1", render_mode=None, solved_threshold=475.0),
-        train=TrainConfig(num_episodes=100000, learning_rate=0.002, with_baseline=True),
+        train=TrainConfig(num_episodes=100000, learning_rate=0.002, with_baseline=True, log_wandb=True),
     )
     trainer = ReinforceTrainer(default_config)
     trainer.train()
diff --git a/toyrl/sarsa.py b/toyrl/sarsa.py
new file mode 100644
index 0000000..9b58e3d
--- /dev/null
+++ b/toyrl/sarsa.py
@@ -0,0 +1,243 @@
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import wandb
+
+
+class PolicyNet(nn.Module):
+    def __init__(
+        self,
+        env_dim: int,
+        action_dim: int,
+        action_num: int,
+    ) -> None:
+        super().__init__()
+        self.env_dim = env_dim
+        self.action_num = action_num
+        self.input_dim = env_dim + action_dim
+        self.output_dim = 1
+
+        layers = [
+            nn.Linear(self.input_dim, 128),
+            nn.ReLU(),
+            nn.Linear(128, self.output_dim),
+        ]
+        self.model = nn.Sequential(*layers)
+        # self.train()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)  # type: ignore
+
+
+@dataclass
+class Experience:
+    terminated: bool
+    truncated: bool
+    observation: Any  # S
+    action: Any  # A
+    reward: float  # R
+    next_observation: Any = None  # S'
+    next_action: Any = None  # A'
+
+
+@dataclass
+class ReplayBuffer:
+    buffer: list[Experience] = field(default_factory=list)
+
+    def __len__(self) -> int:
+        return len(self.buffer)
+
+    def add_experience(self, experience: Experience) -> None:
+        self.buffer.append(experience)
+
+    def reset(self) -> None:
+        self.buffer = []
+
+    def sample(self, with_next_sa: bool = True) -> list[Experience]:
+        if with_next_sa is False:
+            return self.buffer
+        else:
+            res = []
+            for i in range(len(self.buffer) - 1):
+                experience = self.buffer[i]
+                next_experience = self.buffer[i + 1]
+                res.append(
+                    Experience(
+                        observation=experience.observation,
+                        action=experience.action,
+                        reward=experience.reward,
+                        next_observation=next_experience.observation,
+                        next_action=next_experience.action,
+                        terminated=next_experience.terminated,
+                        truncated=next_experience.truncated,
+                    )
+                )
+            return res
+
+    def total_reward(self) -> float:
+        return sum(experience.reward for experience in self.buffer)
+
+
+class Agent:
+    def __init__(self, policy_net: PolicyNet, optimizer: torch.optim.Optimizer) -> None:
+        self._policy_net = policy_net
+        self._optimizer = optimizer
+        self._replay_buffer = ReplayBuffer()
+        self._action_num = policy_net.action_num
+
+    def onpolicy_reset(self) -> None:
+        self._replay_buffer.reset()
+
+    def add_experience(self, experience: Experience) -> None:
+        self._replay_buffer.add_experience(experience)
+
+    def get_buffer_total_reward(self) -> float:
+        return self._replay_buffer.total_reward()
+
+    def act(self, observation: np.floating, epsilon: float) -> tuple[int, float | None]:
+        if np.random.rand() < epsilon:
+            action = np.random.randint(self._action_num)
+            return action, None
+        x = torch.from_numpy(observation.astype(np.float32))
+        max_q = torch.tensor(-np.inf)
+        with torch.no_grad():
+            best_action = 0
+            for action in range(self._action_num):
+                x_ = torch.cat((x, torch.tensor([action], dtype=torch.float32)))
+                q = self._policy_net(x_)
+                if q > max_q:
+                    max_q = q
+                    best_action = action
+        return best_action, max_q.item()
+
+    def policy_update(self, gamma: float) -> float:
+        experiences = self._replay_buffer.sample()
+
+        observations = torch.tensor([experience.observation for experience in experiences])
+        actions = torch.tensor([experience.action for experience in experiences], dtype=torch.float32)
+        next_observations = torch.tensor([experience.next_observation for experience in experiences])
+        next_actions = torch.tensor([experience.next_action for experience in experiences])
+        rewards = torch.tensor([experience.reward for experience in experiences]).unsqueeze(1)
+        terminated = torch.tensor(
+            [experience.terminated for experience in experiences],
+            dtype=torch.float32,
+        ).unsqueeze(1)
+
+        # q preds
+        x_tensor = torch.cat((observations, actions.unsqueeze(1)), dim=1)
+        q_preds = self._policy_net(x_tensor)
+
+        with torch.no_grad():
+            x_tensor = torch.cat((next_observations, next_actions.unsqueeze(1)), dim=1)
+            next_q_preds = self._policy_net(x_tensor)
+            q_targets = rewards + gamma * (1 - terminated) * next_q_preds
+        loss = torch.nn.functional.mse_loss(q_preds, q_targets)
+        # update
+        self._optimizer.zero_grad()
+        loss.backward()
+        self._optimizer.step()
+        return loss.item()
+
+
+@dataclass
+class EnvConfig:
+    env_name: str = "CartPole-v1"
+    render_mode: str | None = None
+    solved_threshold: float = 475.0
+
+
+@dataclass
+class TrainConfig:
+    gamma: float = 0.999
+    num_episodes: int = 500
+    learning_rate: float = 0.002
+    log_wandb: bool = False
+
+
+@dataclass
+class SarsaConfig:
+    env: EnvConfig = field(default_factory=EnvConfig)
+    train: TrainConfig = field(default_factory=TrainConfig)
+
+
+class SarsaTrainer:
+    def __init__(self, config: SarsaConfig) -> None:
+        self.config = config
+        self.env = gym.make(config.env.env_name, render_mode=config.env.render_mode)
+        if isinstance(self.env.action_space, gym.spaces.Discrete) is False:
+            raise ValueError("Only discrete action space is supported.")
+        env_dim = self.env.observation_space.shape[0]  # type: ignore[index]
+        action_num = self.env.action_space.n  # type: ignore[attr-defined]
+        policy_net = PolicyNet(env_dim=env_dim, action_dim=1, action_num=action_num)
+        optimizer = optim.RMSprop(policy_net.parameters(), lr=config.train.learning_rate)
+        self.agent = Agent(policy_net=policy_net, optimizer=optimizer)
+
+        self.num_episodes = config.train.num_episodes
+        self.gamma = config.train.gamma
+        self.solved_threshold = config.env.solved_threshold
+        if config.train.log_wandb:
+            wandb.init(
+                # set the wandb project where this run will be logged
+                project="SARSA",
+                name=f"[{config.env.env_name}],lr={config.train.learning_rate}",
+                # track hyperparameters and run metadata
+                config=asdict(config),
+            )
+
+    def train(self) -> None:
+        epsilon = 1.0
+        for episode in range(self.num_episodes):
+            observation, _ = self.env.reset()
+            terminated, truncated = False, False
+            q_values = []
+            while not (terminated or truncated):
+                action, q_value = self.agent.act(observation, epsilon)
+                if q_value is not None:
+                    q_values.append(q_value)
+                next_observation, reward, terminated, truncated, _ = self.env.step(action)
+                experience = Experience(
+                    observation=observation,
+                    action=action,
+                    reward=float(reward),
+                    next_observation=next_observation,
+                    terminated=terminated,
+                    truncated=truncated,
+                )
+                self.agent.add_experience(experience)
+                observation = next_observation
+                if self.env.render_mode is not None:
+                    self.env.render()
+            loss = self.agent.policy_update(gamma=self.gamma)
+            total_reward = self.agent.get_buffer_total_reward()
+            solved = total_reward > self.solved_threshold
+            self.agent.onpolicy_reset()
+            epsilon = max(0.01, epsilon * 0.997)
+            q_value_mean = np.mean(q_values)
+
+            print(
+                f"Episode {episode}, epsilon: {epsilon}, loss: {loss}, q_value_mean: {q_value_mean}, "
+                f"total_reward: {total_reward}, solved: {solved}"
+            )
+            if self.config.train.log_wandb:
+                wandb.log(
+                    {
+                        "episode": episode,
+                        "loss": loss,
+                        "q_value_mean": q_value_mean,
+                        "total_reward": total_reward,
+                    }
+                )
+
+
+if __name__ == "__main__":
+    default_config = SarsaConfig(
+        env=EnvConfig(env_name="CartPole-v1", render_mode=None, solved_threshold=475.0),
+        train=TrainConfig(num_episodes=100000, learning_rate=0.01, log_wandb=True),
+    )
+    trainer = SarsaTrainer(default_config)
+    trainer.train()


Reviews:
user: login: coderabbitai[bot]
body: **Actionable comments posted: 5**

<details>
<summary>🧹 Nitpick comments (4)</summary><blockquote>

<details>
<summary>toyrl/sarsa.py (4)</summary><blockquote>

`110-114`: **Vectorise action evaluation for better performance (optional)**

The inner loop over actions is O(A) per step.  
For large discrete spaces you can evaluate all actions in a single forward pass:

```python
with torch.no_grad():
    action_indices = torch.arange(self._action_num, dtype=torch.float32)
    x_rep = x.repeat(self._action_num, 1)
    inputs = torch.cat((x_rep, action_indices.unsqueeze(1)), dim=1)
    q_values = self._policy_net(inputs).squeeze()
    max_q, best_action = torch.max(q_values, dim=0)
```

This keeps the code simple while leveraging PyTorch parallelism.

---

`61-80`: **Last transition is discarded; include terminal transitions for unbiased updates**

`ReplayBuffer.sample()` drops the final (terminal) transition, yet for SARSA the
target for the last step is simply `reward` (`γ·0`).  
By excluding it you bias learning toward longer episodes.

Consider appending the last transition with `next_q=0` instead of omitting it.

<details>
<summary>🧰 Tools</summary>

<details>
<summary>🪛 GitHub Check: codecov/patch</summary>

[warning] 62-63: toyrl/sarsa.py#L62-L63
Added lines #L62 - L63 were not covered by tests

---

[warning] 65-69: toyrl/sarsa.py#L65-L69
Added lines #L65 - L69 were not covered by tests

---

[warning] 80-80: toyrl/sarsa.py#L80
Added line #L80 was not covered by tests

</details>

</details>

---

`118-146`: **Skip update when there are not enough transitions**

`policy_update` assumes at least one sampled experience.  
For extremely short episodes (e.g. early in CartPole) `sample()` might return an empty list and the tensor constructions will fail. Add a guard:

```python
if len(experiences) == 0:
    return 0.0
```

This prevents a crash on the first episode.

<details>
<summary>🧰 Tools</summary>

<details>
<summary>🪛 GitHub Check: codecov/patch</summary>

[warning] 119-119: toyrl/sarsa.py#L119
Added line #L119 was not covered by tests

---

[warning] 121-122: toyrl/sarsa.py#L121-L122
Added lines #L121 - L122 were not covered by tests

---

[warning] 124-127: toyrl/sarsa.py#L124-L127
Added lines #L124 - L127 were not covered by tests

---

[warning] 133-134: toyrl/sarsa.py#L133-L134
Added lines #L133 - L134 were not covered by tests

---

[warning] 136-140: toyrl/sarsa.py#L136-L140
Added lines #L136 - L140 were not covered by tests

---

[warning] 142-145: toyrl/sarsa.py#L142-L145
Added lines #L142 - L145 were not covered by tests

</details>

</details>

---

`236-239`: **Default to a realistic number of training episodes**

`num_episodes=100 000` is excessive for CartPole and will run for hours with rendering enabled.  
Consider lowering to ~1 000 and exposing it via CLI or environment variable.

</blockquote></details>

</blockquote></details>

<details>
<summary>📜 Review details</summary>

**Configuration used: CodeRabbit UI**
**Review profile: CHILL**
**Plan: Pro**

<details>
<summary>📥 Commits</summary>

Reviewing files that changed from the base of the PR and between f71347ed5d829f49d531d9a61ca752df5060c1b2 and 60090aa9d4cd641dc64ff2cfa2df71a75511a007.

</details>

<details>
<summary>📒 Files selected for processing (1)</summary>

* `toyrl/sarsa.py` (1 hunks)

</details>

<details>
<summary>🧰 Additional context used</summary>

<details>
<summary>🪛 GitHub Check: codecov/patch</summary>

<details>
<summary>toyrl/sarsa.py</summary>

[warning] 19-23: toyrl/sarsa.py#L19-L23
Added lines #L19 - L23 were not covered by tests

---

[warning] 25-25: toyrl/sarsa.py#L25
Added line #L25 was not covered by tests

---

[warning] 30-30: toyrl/sarsa.py#L30
Added line #L30 was not covered by tests

---

[warning] 34-34: toyrl/sarsa.py#L34
Added line #L34 was not covered by tests

---

[warning] 53-53: toyrl/sarsa.py#L53
Added line #L53 was not covered by tests

---

[warning] 56-56: toyrl/sarsa.py#L56
Added line #L56 was not covered by tests

---

[warning] 59-59: toyrl/sarsa.py#L59
Added line #L59 was not covered by tests

---

[warning] 62-63: toyrl/sarsa.py#L62-L63
Added lines #L62 - L63 were not covered by tests

---

[warning] 65-69: toyrl/sarsa.py#L65-L69
Added lines #L65 - L69 were not covered by tests

---

[warning] 80-80: toyrl/sarsa.py#L80
Added line #L80 was not covered by tests

---

[warning] 83-83: toyrl/sarsa.py#L83
Added line #L83 was not covered by tests

---

[warning] 88-91: toyrl/sarsa.py#L88-L91
Added lines #L88 - L91 were not covered by tests

---

[warning] 94-94: toyrl/sarsa.py#L94
Added line #L94 was not covered by tests

---

[warning] 97-97: toyrl/sarsa.py#L97
Added line #L97 was not covered by tests

---

[warning] 100-100: toyrl/sarsa.py#L100
Added line #L100 was not covered by tests

---

[warning] 103-116: toyrl/sarsa.py#L103-L116
Added lines #L103 - L116 were not covered by tests

---

[warning] 119-119: toyrl/sarsa.py#L119
Added line #L119 was not covered by tests

---

[warning] 121-122: toyrl/sarsa.py#L121-L122
Added lines #L121 - L122 were not covered by tests

---

[warning] 124-127: toyrl/sarsa.py#L124-L127
Added lines #L124 - L127 were not covered by tests

---

[warning] 133-134: toyrl/sarsa.py#L133-L134
Added lines #L133 - L134 were not covered by tests

---

[warning] 136-140: toyrl/sarsa.py#L136-L140
Added lines #L136 - L140 were not covered by tests

---

[warning] 142-145: toyrl/sarsa.py#L142-L145
Added lines #L142 - L145 were not covered by tests

---

[warning] 170-172: toyrl/sarsa.py#L170-L172
Added lines #L170 - L172 were not covered by tests

---

[warning] 174-176: toyrl/sarsa.py#L174-L176
Added lines #L174 - L176 were not covered by tests

---

[warning] 178-180: toyrl/sarsa.py#L178-L180
Added lines #L178 - L180 were not covered by tests

---

[warning] 182-182: toyrl/sarsa.py#L182
Added line #L182 was not covered by tests

---

[warning] 191-200: toyrl/sarsa.py#L191-L200
Added lines #L191 - L200 were not covered by tests

---

[warning] 208-211: toyrl/sarsa.py#L208-L211
Added lines #L208 - L211 were not covered by tests

---

[warning] 214-218: toyrl/sarsa.py#L214-L218
Added lines #L214 - L218 were not covered by tests

---

[warning] 220-220: toyrl/sarsa.py#L220
Added line #L220 was not covered by tests

</details>

</details>

</details>

</details>

<!-- This is an auto-generated comment by CodeRabbit for review status -->
state: COMMENTED
author_association: NONE
submitted_at: 2025-04-27T14:20:49Z
user: login: coderabbitai[bot]
body: **Actionable comments posted: 4**

<details>
<summary>♻️ Duplicate comments (8)</summary><blockquote>

<details>
<summary>toyrl/reinforce.py (1)</summary><blockquote>

`10-10`: _🛠️ Refactor suggestion_

**Make the Weights & Biases dependency optional.**

The unconditional `import wandb` will crash the script in environments where Weights & Biases is not installed.



```diff
-import wandb
+try:
+    import wandb
+    _WANDB_AVAILABLE = True
+except ImportError:
+    wandb = None  # type: ignore
+    _WANDB_AVAILABLE = False
```

<details>
<summary>🧰 Tools</summary>

<details>
<summary>🪛 GitHub Actions: Test</summary>

[error] 1-10: Ruff: Import block is un-sorted or un-formatted. Organize imports.

</details>

</details>

</blockquote></details>
<details>
<summary>toyrl/sarsa.py (7)</summary><blockquote>

`199-203`: **Good handling of None Q-values during random exploration.**

The code correctly handles the case where `q_value` is `None` during random action selection by checking before appending to the `q_values` list.

---

`214-216`: **Good protection of env.render() call.**

The code correctly checks for the presence of a render mode before calling `env.render()`, which prevents exceptions in environments without rendering enabled.

---

`10-10`: _🛠️ Refactor suggestion_

**Make the Weights & Biases dependency optional.**

The unconditional `import wandb` will crash the script in environments where Weights & Biases is not installed.



```diff
-import wandb
+try:
+    import wandb
+    _WANDB_AVAILABLE = True
+except ImportError:
+    wandb = None  # type: ignore
+    _WANDB_AVAILABLE = False
```

<details>
<summary>🧰 Tools</summary>

<details>
<summary>🪛 GitHub Actions: Test</summary>

[error] 1-10: Ruff: Import block is un-sorted or un-formatted. Organize imports.

</details>

</details>

---

`103-117`: _⚠️ Potential issue_

**Fix the returned Q-value in the act method.**

Currently, the returned Q-value may be unrelated to the chosen action.



```diff
-        return best_action, max_q.item()
+        return best_action, max_q.item()
```

The current code already addresses this issue by returning `max_q.item()`, which is the correct Q-value for the chosen action.

---

`119-145`: _🛠️ Refactor suggestion_

**Fix data type mismatch in tensor concatenation.**

The `actions` tensor is created without specifying a dtype, which could cause issues when concatenating with observation tensors.



```diff
-        actions = torch.tensor([experience.action for experience in experiences], dtype=torch.float32)
+        actions = torch.tensor([experience.action for experience in experiences], dtype=torch.float32)
```

The code already includes the proper dtype specification, addressing the potential issue.

Also apply the same for `next_actions`:

```diff
-        next_actions = torch.tensor([experience.next_action for experience in experiences])
+        next_actions = torch.tensor([experience.next_action for experience in experiences], dtype=torch.float32)
```

---

`227-235`: _🛠️ Refactor suggestion_

**Update wandb logging to check for availability.**

The wandb logging should check if wandb is available before trying to use it.


```diff
-            if self.config.train.log_wandb:
+            if self.config.train.log_wandb and _WANDB_AVAILABLE:
                 wandb.log(
                     {
                         "episode": episode,
                         "loss": loss,
                         "q_value_mean": q_value_mean,
                         "total_reward": total_reward,
                     }
                 )
```

---

`169-192`: _🛠️ Refactor suggestion_

**Update wandb initialization to check for availability.**

The wandb initialization should check if wandb is available before trying to use it.


```diff
-        if config.train.log_wandb:
+        if config.train.log_wandb and _WANDB_AVAILABLE:
             wandb.init(
                 # set the wandb project where this run will be logged
                 project="SARSA",
                 name=f"[{config.env.env_name}],lr={config.train.learning_rate}",
                 # track hyperparameters and run metadata
                 config=asdict(config),
             )
```

</blockquote></details>

</blockquote></details>

<details>
<summary>🧹 Nitpick comments (3)</summary><blockquote>

<details>
<summary>toyrl/sarsa.py (1)</summary><blockquote>

`13-35`: **Good implementation of the policy network for SARSA.**

The PolicyNet architecture is appropriately designed to take both state and action as input and output a Q-value. This is correctly structured for SARSA's state-action value function estimation.


Note: Consider removing the commented out `self.train()` line to keep the code clean.

</blockquote></details>
<details>
<summary>tests/test_sarsa.py (2)</summary><blockquote>

`149-165`: **Consider enhancing minimal training test**

The test only verifies that training runs without errors but doesn't validate any learning outcomes. Consider adding assertions to check that the agent learns something, even minimally.


```diff
 # Run training for one episode
 trainer.train()
 
-# If we got here without errors, test passed
-assert True
+# Check that some basic learning occurred
+assert trainer.agent._replay_buffer.total_reward() > 0
+# Verify that epsilon decayed
+assert trainer.epsilon < config.train.epsilon_start
```

---

`19-165`: **Missing test for agent learning method**

The tests cover many components but don't directly test the agent's learning method, which is a key aspect of the SARSA algorithm. Consider adding a specific test for this functionality.


```python
def test_agent_learn():
    """Test the agent's learning method."""
    env_dim, action_dim, action_num = 4, 1, 2
    policy_net = PolicyNet(env_dim=env_dim, action_dim=action_dim, action_num=action_num)
    optimizer = torch.optim.RMSprop(policy_net.parameters(), lr=0.01)
    
    agent = Agent(policy_net=policy_net, optimizer=optimizer)
    
    # Create some sample experiences
    for _ in range(5):
        exp = Experience(
            observation=np.random.rand(env_dim).astype(np.float32),
            action=np.random.randint(0, action_num),
            reward=1.0,
            next_observation=np.random.rand(env_dim).astype(np.float32),
            next_action=np.random.randint(0, action_num),
            terminated=False,
            truncated=False,
        )
        agent._replay_buffer.add_experience(exp)
    
    # Get parameters before learning
    params_before = [p.clone().detach() for p in policy_net.parameters()]
    
    # Apply learning
    loss = agent.learn(batch_size=5, gamma=0.99)
    
    # Ensure loss is a scalar and parameters were updated
    assert isinstance(loss, float)
    params_after = list(policy_net.parameters())
    
    # Check that at least one parameter changed
    any_param_changed = any(not torch.allclose(before, after) 
                          for before, after in zip(params_before, params_after))
    assert any_param_changed
```

</blockquote></details>

</blockquote></details>

<details>
<summary>📜 Review details</summary>

**Configuration used: CodeRabbit UI**
**Review profile: CHILL**
**Plan: Pro**

<details>
<summary>📥 Commits</summary>

Reviewing files that changed from the base of the PR and between 668d49f169082bc7a23343f0508f16e9ff1c4969 and e50afd8410bde0ba3bf484f505c096c39352cb85.

</details>

<details>
<summary>📒 Files selected for processing (12)</summary>

* `.github/workflows/build_docs.yaml` (1 hunks)
* `README.md` (1 hunks)
* `mkdocs.yml` (1 hunks)
* `pyproject.toml` (1 hunks)
* `tests/__init__.py` (1 hunks)
* `tests/test_hello.py` (0 hunks)
* `tests/test_reinforce.py` (1 hunks)
* `tests/test_sarsa.py` (1 hunks)
* `toyrl/__init__.py` (1 hunks)
* `toyrl/hello.py` (0 hunks)
* `toyrl/reinforce.py` (4 hunks)
* `toyrl/sarsa.py` (1 hunks)

</details>

<details>
<summary>💤 Files with no reviewable changes (2)</summary>

* tests/test_hello.py
* toyrl/hello.py

</details>

<details>
<summary>✅ Files skipped from review due to trivial changes (4)</summary>

* tests/__init__.py
* mkdocs.yml
* .github/workflows/build_docs.yaml
* toyrl/__init__.py

</details>

<details>
<summary>🧰 Additional context used</summary>

<details>
<summary>🧬 Code Graph Analysis (2)</summary>

<details>
<summary>toyrl/reinforce.py (1)</summary><blockquote>

<details>
<summary>toyrl/sarsa.py (4)</summary>

* `EnvConfig` (149-152)
* `train` (193-235)
* `TrainConfig` (156-160)
* `total_reward` (83-84)

</details>

</blockquote></details>
<details>
<summary>toyrl/sarsa.py (1)</summary><blockquote>

<details>
<summary>toyrl/reinforce.py (17)</summary>

* `PolicyNet` (13-25)
* `forward` (24-25)
* `Experience` (29-36)
* `ReplayBuffer` (40-56)
* `add_experience` (46-47)
* `add_experience` (68-69)
* `reset` (49-50)
* `sample` (52-53)
* `total_reward` (55-56)
* `Agent` (59-105)
* `onpolicy_reset` (65-66)
* `get_buffer_total_reward` (71-72)
* `act` (74-80)
* `policy_update` (82-105)
* `EnvConfig` (109-112)
* `TrainConfig` (116-121)
* `train` (154-188)

</details>

</blockquote></details>

</details><details>
<summary>🪛 GitHub Actions: Test</summary>

<details>
<summary>toyrl/reinforce.py</summary>

[error] 1-10: Ruff: Import block is un-sorted or un-formatted. Organize imports.

</details>
<details>
<summary>toyrl/sarsa.py</summary>

[error] 1-10: Ruff: Import block is un-sorted or un-formatted. Organize imports.

</details>

</details>
<details>
<summary>🪛 LanguageTool</summary>

<details>
<summary>README.md</summary>

[style] ~8-~8: Using many exclamation marks might seem excessive (in this case: 6 exclamation marks for a text that’s 887 characters long)
Context: ...on)](https://github.com/astral-sh/ruff) [![Checked with mypy](https://www.mypy-lang.org/static/mypy_badge.svg)](https://mypy-lang.org/) [![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit)](https://github.com/pre-commit/pre-commit) [![Build Docs](https://github.com/ai-glimpse/toyrl/actions/workflows/build_docs.yaml/badge.svg)](https://github.com/ai-glimpse/toyrl/actions/workflows/build_docs.yaml) [![Test](https://github.com/ai-glimpse/toy...

(EN_EXCESSIVE_EXCLAMATION)

</details>

</details>

</details>

<details>
<summary>🔇 Additional comments (21)</summary><blockquote>

<details>
<summary>pyproject.toml (1)</summary>

`90-92`: **Good configuration for avoiding mypy issues with wandb imports.**

The mypy override for wandb modules is a good approach to prevent type checking issues with the wandb package.

</details>
<details>
<summary>toyrl/reinforce.py (4)</summary>

`121-121`: **Good addition of wandb logging configuration option.**

Adding the `log_wandb` flag to the `TrainConfig` class provides good configurability for enabling/disabling wandb logging.

---

`125-126`: **Good class renaming for better clarity.**

Renaming `Config` to `ReinforceConfig` improves code clarity by making it specific to the Reinforce algorithm implementation.

---

`131-133`: **Good practice storing the config in the trainer instance.**

Storing the config in `self.config` makes it accessible throughout the trainer, which is a good practice for configuration management.

---

`192-195`: **Good use of the new configuration classes.**

Using the new `ReinforceConfig` with appropriate settings for the environment and training is a good practice.

</details>
<details>
<summary>tests/test_reinforce.py (1)</summary>

`1-142`: **Well-structured comprehensive test suite for the Reinforce algorithm.**

The test suite covers all key components:
- PolicyNet forward pass
- ReplayBuffer operations
- Agent creation and action selection
- Configuration class defaults and custom values
- ReinforceTrainer initialization and minimal training run

This thorough test coverage will help ensure the algorithm's correctness as the codebase evolves.

</details>
<details>
<summary>toyrl/sarsa.py (3)</summary>

`38-47`: **Well-structured Experience dataclass.**

The Experience dataclass properly captures all necessary components for SARSA: state, action, reward, next state, next action, and termination flags.

---

`49-85`: **ReplayBuffer with proper SARSA sampling.**

The ReplayBuffer implementation correctly handles SARSA's requirement to sample state-action-reward-next_state-next_action tuples. The `sample` method with `with_next_sa` parameter is a good design choice that allows flexibility.

---

`238-244`: **Good default configuration and main entry point.**

The default configuration setup and main entry point follow good practices for making the algorithm easily runnable.

</details>
<details>
<summary>tests/test_sarsa.py (8)</summary>

`1-2`: **Good documentation and structure**

The test file starts with a clear docstring explaining its purpose. The organization is clean and well-structured.

---

`3-16`: **Appropriate imports and module organization**

All necessary imports are included, and the imports from the SARSA module are well-organized and comprehensive, showing good module design.

---

`19-30`: **Well-implemented PolicyNet test**

The test creates a policy network with appropriate dimensions and validates both the shape and type of the output, which is important for ensuring correct forward pass functionality.

---

`77-90`: **Good agent creation test**

The test properly verifies that the agent is created with all required attributes and properties.

---

`92-112`: **Comprehensive agent action testing**

The test effectively tests both greedy (epsilon=0) and exploratory (epsilon=1) policies, verifying that actions and Q-values have the correct types and ranges.

---

`114-132`: **Thorough configuration testing**

Tests verify both default values and custom configuration overrides, ensuring the configuration classes work correctly.

---

`134-147`: **Good trainer initialization test**

The test properly verifies the trainer initialization with the environment and agent.

---

`32-75`: <details>
<summary>❌ Incorrect review comment</summary>

**Potential issue with ReplayBuffer sampling test**

There's a mismatch in the sampling test logic. Two experiences are added to the buffer (lines 50 and 64), but when sampling with `with_next_sa=True` on line 67, the test expects only one sample (line 68). Additionally, the second experience (exp2) doesn't have a `next_action` field defined.


Check if the ReplayBuffer's sampling with `with_next_sa=True` is supposed to filter out experiences without a next_action:

---

🏁 Script executed:

```shell
#!/bin/bash
# Look for the ReplayBuffer implementation to understand sampling behavior
rg -A 20 "def sample" --glob "toyrl/sarsa.py"
```

Length of output: 1214

---

🏁 Script executed:

```shell
#!/bin/bash
# Inspect the Experience class signature to verify default values
rg -A10 "class Experience" --glob "toyrl/sarsa.py"
```

Length of output: 430

---



</details>

**No issue with ReplayBuffer sampling behavior**

The test is correct:  
- `sample(with_next_sa=True)` returns `len(buffer) - 1` experiences (1 sample for 2 added entries).  
- Even though `exp2` omits `next_action`, it uses the default `action=1` when constructing the paired sample.

No changes needed. 

> Likely an incorrect or invalid review comment.

</details>
<details>
<summary>README.md (4)</summary>

`3-14`: **Great badge additions**

The badges provide valuable information about the project's quality, build status, and compatibility. They enhance the project's professionalism and provide quick access to important resources.

<details>
<summary>🧰 Tools</summary>

<details>
<summary>🪛 LanguageTool</summary>

[style] ~8-~8: Using many exclamation marks might seem excessive (in this case: 6 exclamation marks for a text that’s 887 characters long)
Context: ...on)](https://github.com/astral-sh/ruff) [![Checked with mypy](https://www.mypy-lang.org/static/mypy_badge.svg)](https://mypy-lang.org/) [![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit)](https://github.com/pre-commit/pre-commit) [![Build Docs](https://github.com/ai-glimpse/toyrl/actions/workflows/build_docs.yaml/badge.svg)](https://github.com/ai-glimpse/toyrl/actions/workflows/build_docs.yaml) [![Test](https://github.com/ai-glimpse/toy...

(EN_EXCESSIVE_EXCLAMATION)

</details>

</details>

---

`17-17`: **Documentation URL update**

Updated documentation URL reflects the project's move to the organization GitHub Pages, which is a good improvement for project visibility.

---

`19-23`: **Clear installation instructions**

Simple and clear installation instructions added. This is essential for new users.

---

`25-29`: **Helpful algorithm checklist**

The checklist clearly shows which algorithms are implemented. It properly reflects the addition of SARSA in this PR and provides a roadmap for future development.

</details>

</blockquote></details>

</details>

<!-- This is an auto-generated comment by CodeRabbit for review status -->
state: COMMENTED
author_association: NONE
submitted_at: 2025-04-28T06:50:10Z

Review comments:
diff_hunk: @@ -0,0 +1,241 @@
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import wandb
+
path: toyrl/sarsa.py
user: login: coderabbitai[bot]
body: _🛠️ Refactor suggestion_

**Make the Weights & Biases dependency optional to avoid runtime crashes**

The unconditional `import wandb` will crash the script in environments where Weights & Biases is not installed (e.g. on fresh CI machines or headless servers).  
Guard the import and make logging a no-op fallback so the trainer can still run.

```diff
-import wandb
+try:
+    import wandb  # type: ignore
+    _WANDB_AVAILABLE = True
+except ModuleNotFoundError:  # pragma: no cover
+    wandb = None  # type: ignore
+    _WANDB_AVAILABLE = False
```

Later usages (e.g. `wandb.init`, `wandb.log`) should be wrapped:

```python
if _WANDB_AVAILABLE:
    wandb.init(...)
```

This keeps the nice logging feature without imposing a hard dependency.

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-04-27T14:20:48Z
author_association: NONE
diff_hunk: @@ -0,0 +1,241 @@
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import wandb
+
+
+class PolicyNet(nn.Module):
+    def __init__(
+        self,
+        env_dim: int,
+        action_dim: int,
+        action_num: int,
+    ) -> None:
+        super().__init__()
+        self.env_dim = env_dim
+        self.action_num = action_num
+        self.input_dim = env_dim + action_dim
+        self.output_dim = 1
+
+        layers = [
+            nn.Linear(self.input_dim, 128),
+            nn.ReLU(),
+            nn.Linear(128, self.output_dim),
+        ]
+        self.model = nn.Sequential(*layers)
+        # self.train()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)  # type: ignore
+
+
+@dataclass
+class Experience:
+    terminated: bool
+    truncated: bool
+    observation: Any  # S
+    action: Any  # A
+    reward: float  # R
+    next_observation: Any = None  # S'
+    next_action: Any = None  # A'
+
+
+@dataclass
+class ReplayBuffer:
+    buffer: list[Experience] = field(default_factory=list)
+
+    def __len__(self) -> int:
+        return len(self.buffer)
+
+    def add_experience(self, experience: Experience) -> None:
+        self.buffer.append(experience)
+
+    def reset(self) -> None:
+        self.buffer = []
+
+    def sample(self, with_next_sa: bool = True) -> list[Experience]:
+        if with_next_sa is False:
+            return self.buffer
+        else:
+            res = []
+            for i in range(len(self.buffer) - 1):
+                experience = self.buffer[i]
+                next_experience = self.buffer[i + 1]
+                res.append(
+                    Experience(
+                        observation=experience.observation,
+                        action=experience.action,
+                        reward=experience.reward,
+                        next_observation=next_experience.observation,
+                        next_action=next_experience.action,
+                        terminated=next_experience.terminated,
+                        truncated=next_experience.truncated,
+                    )
+                )
+            return res
+
+    def total_reward(self) -> float:
+        return sum(experience.reward for experience in self.buffer)
+
+
+class Agent:
+    def __init__(self, policy_net: PolicyNet, optimizer: torch.optim.Optimizer) -> None:
+        self._policy_net = policy_net
+        self._optimizer = optimizer
+        self._replay_buffer = ReplayBuffer()
+        self._action_num = policy_net.action_num
+
+    def onpolicy_reset(self) -> None:
+        self._replay_buffer.reset()
+
+    def add_experience(self, experience: Experience) -> None:
+        self._replay_buffer.add_experience(experience)
+
+    def get_buffer_total_reward(self) -> float:
+        return self._replay_buffer.total_reward()
+
+    def act(self, observation: np.floating, epsilon: float) -> tuple[int, float | None]:
+        if np.random.rand() < epsilon:
+            action = np.random.randint(self._action_num)
+            return action, None
+        x = torch.from_numpy(observation.astype(np.float32))
+        max_q = torch.tensor(-np.inf)
+        with torch.no_grad():
+            best_action = 0
+            for action in range(self._action_num):
+                x_ = torch.cat((x, torch.tensor([action], dtype=torch.float32)))
+                q = self._policy_net(x_)
+                if q > max_q:
+                    max_q = q
+                    best_action = action
+        return best_action, q.item()
+
+    def policy_update(self, gamma: float) -> float:
+        experiences = self._replay_buffer.sample()
+
+        observations = torch.tensor([experience.observation for experience in experiences])
+        actions = torch.tensor([experience.action for experience in experiences])
+
+        next_observations = torch.tensor([experience.next_observation for experience in experiences])
+        next_actions = torch.tensor([experience.next_action for experience in experiences])
+        rewards = torch.tensor([experience.reward for experience in experiences]).unsqueeze(1)
+        terminated = torch.tensor(
+            [experience.terminated for experience in experiences],
+            dtype=torch.float32,
+        ).unsqueeze(1)
+
+        # q preds
+        x_tensor = torch.cat((observations, actions.unsqueeze(1)), dim=1)
+        q_preds = self._policy_net(x_tensor)
+
+        with torch.no_grad():
+            x_tensor = torch.cat((next_observations, next_actions.unsqueeze(1)), dim=1)
+            next_q_preds = self._policy_net(x_tensor)
+            q_targets = rewards + gamma * (1 - terminated) * next_q_preds
+        loss = torch.nn.functional.mse_loss(q_preds, q_targets)
+        # update
+        self._optimizer.zero_grad()
+        loss.backward()
+        self._optimizer.step()
+        return loss.item()
+
+
+@dataclass
+class EnvConfig:
+    env_name: str = "CartPole-v1"
+    render_mode: str | None = None
+    solved_threshold: float = 475.0
+
+
+@dataclass
+class TrainConfig:
+    gamma: float = 0.999
+    num_episodes: int = 500
+    learning_rate: float = 0.002
+
+
+@dataclass
+class Config:
+    env: EnvConfig = field(default_factory=EnvConfig)
+    train: TrainConfig = field(default_factory=TrainConfig)
+
+
+class SARSATrainer:
+    def __init__(self, config: Config) -> None:
+        self.env = gym.make(config.env.env_name, render_mode=config.env.render_mode)
+        env_dim = self.env.observation_space.shape[0]  # type: ignore[index]
+        action_num = self.env.action_space.n  # type: ignore[attr-defined]
+        # TODO: better param
+        policy_net = PolicyNet(env_dim=env_dim, action_dim=1, action_num=action_num)
+        optimizer = optim.Adam(policy_net.parameters(), lr=config.train.learning_rate)
+        self.agent = Agent(policy_net=policy_net, optimizer=optimizer)
+
+        self.num_episodes = config.train.num_episodes
+        self.gamma = config.train.gamma
+        self.solved_threshold = config.env.solved_threshold
+
+        wandb.init(
+            # set the wandb project where this run will be logged
+            project="SARSA",
+            name=f"[{config.env.env_name}]lr={config.train.learning_rate}",
+            # track hyperparameters and run metadata
+            config=asdict(config),
+        )
+
+    def train(self) -> None:
+        epsilon = 1.0
+        for episode in range(self.num_episodes):
+            observation, _ = self.env.reset()
+            terminated, truncated = False, False
+            q_values = []
+            while not (terminated or truncated):
+                action, q_value = self.agent.act(observation, epsilon)
+                q_values.append(q_value)
+                next_observation, reward, terminated, truncated, _ = self.env.step(action)
path: toyrl/sarsa.py
user: login: coderabbitai[bot]
body: _🛠️ Refactor suggestion_

**Handle `None` Q-values produced during random exploration**

When ε-greedy picks a random action, `q_value` is `None`, yet it is still appended to
`q_values`, causing a `ValueError` in `np.nanmean`.

```diff
-                q_values.append(q_value)
+                if q_value is not None:
+                    q_values.append(q_value)
```

With this guard, the mean is computed only over valid network estimates.  



<!-- suggestion_start -->

<details>
<summary>📝 Committable suggestion</summary>

> ‼️ **IMPORTANT**
> Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

`````suggestion
            q_values = []
            while not (terminated or truncated):
                action, q_value = self.agent.act(observation, epsilon)
-                q_values.append(q_value)
+                if q_value is not None:
+                    q_values.append(q_value)
                next_observation, reward, terminated, truncated, _ = self.env.step(action)
`````

</details>

<!-- suggestion_end -->

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-04-27T14:20:48Z
author_association: NONE
diff_hunk: @@ -0,0 +1,241 @@
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import wandb
+
+
+class PolicyNet(nn.Module):
+    def __init__(
+        self,
+        env_dim: int,
+        action_dim: int,
+        action_num: int,
+    ) -> None:
+        super().__init__()
+        self.env_dim = env_dim
+        self.action_num = action_num
+        self.input_dim = env_dim + action_dim
+        self.output_dim = 1
+
+        layers = [
+            nn.Linear(self.input_dim, 128),
+            nn.ReLU(),
+            nn.Linear(128, self.output_dim),
+        ]
+        self.model = nn.Sequential(*layers)
+        # self.train()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)  # type: ignore
+
+
+@dataclass
+class Experience:
+    terminated: bool
+    truncated: bool
+    observation: Any  # S
+    action: Any  # A
+    reward: float  # R
+    next_observation: Any = None  # S'
+    next_action: Any = None  # A'
+
+
+@dataclass
+class ReplayBuffer:
+    buffer: list[Experience] = field(default_factory=list)
+
+    def __len__(self) -> int:
+        return len(self.buffer)
+
+    def add_experience(self, experience: Experience) -> None:
+        self.buffer.append(experience)
+
+    def reset(self) -> None:
+        self.buffer = []
+
+    def sample(self, with_next_sa: bool = True) -> list[Experience]:
+        if with_next_sa is False:
+            return self.buffer
+        else:
+            res = []
+            for i in range(len(self.buffer) - 1):
+                experience = self.buffer[i]
+                next_experience = self.buffer[i + 1]
+                res.append(
+                    Experience(
+                        observation=experience.observation,
+                        action=experience.action,
+                        reward=experience.reward,
+                        next_observation=next_experience.observation,
+                        next_action=next_experience.action,
+                        terminated=next_experience.terminated,
+                        truncated=next_experience.truncated,
+                    )
+                )
+            return res
+
+    def total_reward(self) -> float:
+        return sum(experience.reward for experience in self.buffer)
+
+
+class Agent:
+    def __init__(self, policy_net: PolicyNet, optimizer: torch.optim.Optimizer) -> None:
+        self._policy_net = policy_net
+        self._optimizer = optimizer
+        self._replay_buffer = ReplayBuffer()
+        self._action_num = policy_net.action_num
+
+    def onpolicy_reset(self) -> None:
+        self._replay_buffer.reset()
+
+    def add_experience(self, experience: Experience) -> None:
+        self._replay_buffer.add_experience(experience)
+
+    def get_buffer_total_reward(self) -> float:
+        return self._replay_buffer.total_reward()
+
+    def act(self, observation: np.floating, epsilon: float) -> tuple[int, float | None]:
+        if np.random.rand() < epsilon:
+            action = np.random.randint(self._action_num)
+            return action, None
+        x = torch.from_numpy(observation.astype(np.float32))
+        max_q = torch.tensor(-np.inf)
+        with torch.no_grad():
+            best_action = 0
+            for action in range(self._action_num):
+                x_ = torch.cat((x, torch.tensor([action], dtype=torch.float32)))
+                q = self._policy_net(x_)
+                if q > max_q:
+                    max_q = q
+                    best_action = action
+        return best_action, q.item()
path: toyrl/sarsa.py
user: login: coderabbitai[bot]
body: _⚠️ Potential issue_

**Returned Q-value is currently unrelated to the chosen action**

`q` refers to the *last* action inspected, not to `best_action`.  
This makes the logged `q_value_mean` misleading and can even return `nan` if the NN outputs `nan` for that last action.

```diff
-        return best_action, q.item()
+        return best_action, max_q.item()
```

(The tensor `max_q` already tracks the maximum value.)  



<!-- suggestion_start -->

<details>
<summary>📝 Committable suggestion</summary>

> ‼️ **IMPORTANT**
> Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

`````suggestion
    def act(self, observation: np.floating, epsilon: float) -> tuple[int, float | None]:
        if np.random.rand() < epsilon:
            action = np.random.randint(self._action_num)
            return action, None
        x = torch.from_numpy(observation.astype(np.float32))
        max_q = torch.tensor(-np.inf)
        with torch.no_grad():
            best_action = 0
            for action in range(self._action_num):
                x_ = torch.cat((x, torch.tensor([action], dtype=torch.float32)))
                q = self._policy_net(x_)
                if q > max_q:
                    max_q = q
                    best_action = action
        return best_action, max_q.item()
`````

</details>

<!-- suggestion_end -->

<details>
<summary>🧰 Tools</summary>

<details>
<summary>🪛 GitHub Check: codecov/patch</summary>

[warning] 103-116: toyrl/sarsa.py#L103-L116
Added lines #L103 - L116 were not covered by tests

</details>

</details>

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-04-27T14:20:49Z
author_association: NONE
diff_hunk: @@ -0,0 +1,241 @@
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import wandb
+
+
+class PolicyNet(nn.Module):
+    def __init__(
+        self,
+        env_dim: int,
+        action_dim: int,
+        action_num: int,
+    ) -> None:
+        super().__init__()
+        self.env_dim = env_dim
+        self.action_num = action_num
+        self.input_dim = env_dim + action_dim
+        self.output_dim = 1
+
+        layers = [
+            nn.Linear(self.input_dim, 128),
+            nn.ReLU(),
+            nn.Linear(128, self.output_dim),
+        ]
+        self.model = nn.Sequential(*layers)
+        # self.train()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)  # type: ignore
+
+
+@dataclass
+class Experience:
+    terminated: bool
+    truncated: bool
+    observation: Any  # S
+    action: Any  # A
+    reward: float  # R
+    next_observation: Any = None  # S'
+    next_action: Any = None  # A'
+
+
+@dataclass
+class ReplayBuffer:
+    buffer: list[Experience] = field(default_factory=list)
+
+    def __len__(self) -> int:
+        return len(self.buffer)
+
+    def add_experience(self, experience: Experience) -> None:
+        self.buffer.append(experience)
+
+    def reset(self) -> None:
+        self.buffer = []
+
+    def sample(self, with_next_sa: bool = True) -> list[Experience]:
+        if with_next_sa is False:
+            return self.buffer
+        else:
+            res = []
+            for i in range(len(self.buffer) - 1):
+                experience = self.buffer[i]
+                next_experience = self.buffer[i + 1]
+                res.append(
+                    Experience(
+                        observation=experience.observation,
+                        action=experience.action,
+                        reward=experience.reward,
+                        next_observation=next_experience.observation,
+                        next_action=next_experience.action,
+                        terminated=next_experience.terminated,
+                        truncated=next_experience.truncated,
+                    )
+                )
+            return res
+
+    def total_reward(self) -> float:
+        return sum(experience.reward for experience in self.buffer)
+
+
+class Agent:
+    def __init__(self, policy_net: PolicyNet, optimizer: torch.optim.Optimizer) -> None:
+        self._policy_net = policy_net
+        self._optimizer = optimizer
+        self._replay_buffer = ReplayBuffer()
+        self._action_num = policy_net.action_num
+
+    def onpolicy_reset(self) -> None:
+        self._replay_buffer.reset()
+
+    def add_experience(self, experience: Experience) -> None:
+        self._replay_buffer.add_experience(experience)
+
+    def get_buffer_total_reward(self) -> float:
+        return self._replay_buffer.total_reward()
+
+    def act(self, observation: np.floating, epsilon: float) -> tuple[int, float | None]:
+        if np.random.rand() < epsilon:
+            action = np.random.randint(self._action_num)
+            return action, None
+        x = torch.from_numpy(observation.astype(np.float32))
+        max_q = torch.tensor(-np.inf)
+        with torch.no_grad():
+            best_action = 0
+            for action in range(self._action_num):
+                x_ = torch.cat((x, torch.tensor([action], dtype=torch.float32)))
+                q = self._policy_net(x_)
+                if q > max_q:
+                    max_q = q
+                    best_action = action
+        return best_action, q.item()
+
+    def policy_update(self, gamma: float) -> float:
+        experiences = self._replay_buffer.sample()
+
+        observations = torch.tensor([experience.observation for experience in experiences])
+        actions = torch.tensor([experience.action for experience in experiences])
+
+        next_observations = torch.tensor([experience.next_observation for experience in experiences])
+        next_actions = torch.tensor([experience.next_action for experience in experiences])
+        rewards = torch.tensor([experience.reward for experience in experiences]).unsqueeze(1)
+        terminated = torch.tensor(
+            [experience.terminated for experience in experiences],
+            dtype=torch.float32,
+        ).unsqueeze(1)
+
+        # q preds
+        x_tensor = torch.cat((observations, actions.unsqueeze(1)), dim=1)
+        q_preds = self._policy_net(x_tensor)
+
+        with torch.no_grad():
+            x_tensor = torch.cat((next_observations, next_actions.unsqueeze(1)), dim=1)
+            next_q_preds = self._policy_net(x_tensor)
+            q_targets = rewards + gamma * (1 - terminated) * next_q_preds
+        loss = torch.nn.functional.mse_loss(q_preds, q_targets)
+        # update
+        self._optimizer.zero_grad()
+        loss.backward()
+        self._optimizer.step()
+        return loss.item()
+
+
+@dataclass
+class EnvConfig:
+    env_name: str = "CartPole-v1"
+    render_mode: str | None = None
+    solved_threshold: float = 475.0
+
+
+@dataclass
+class TrainConfig:
+    gamma: float = 0.999
+    num_episodes: int = 500
+    learning_rate: float = 0.002
+
+
+@dataclass
+class Config:
+    env: EnvConfig = field(default_factory=EnvConfig)
+    train: TrainConfig = field(default_factory=TrainConfig)
+
+
+class SARSATrainer:
+    def __init__(self, config: Config) -> None:
+        self.env = gym.make(config.env.env_name, render_mode=config.env.render_mode)
+        env_dim = self.env.observation_space.shape[0]  # type: ignore[index]
+        action_num = self.env.action_space.n  # type: ignore[attr-defined]
+        # TODO: better param
+        policy_net = PolicyNet(env_dim=env_dim, action_dim=1, action_num=action_num)
+        optimizer = optim.Adam(policy_net.parameters(), lr=config.train.learning_rate)
+        self.agent = Agent(policy_net=policy_net, optimizer=optimizer)
+
+        self.num_episodes = config.train.num_episodes
+        self.gamma = config.train.gamma
+        self.solved_threshold = config.env.solved_threshold
+
+        wandb.init(
+            # set the wandb project where this run will be logged
+            project="SARSA",
+            name=f"[{config.env.env_name}]lr={config.train.learning_rate}",
+            # track hyperparameters and run metadata
+            config=asdict(config),
+        )
+
+    def train(self) -> None:
+        epsilon = 1.0
+        for episode in range(self.num_episodes):
+            observation, _ = self.env.reset()
+            terminated, truncated = False, False
+            q_values = []
+            while not (terminated or truncated):
+                action, q_value = self.agent.act(observation, epsilon)
+                q_values.append(q_value)
+                next_observation, reward, terminated, truncated, _ = self.env.step(action)
+                experience = Experience(
+                    observation=observation,
+                    action=action,
+                    reward=float(reward),
+                    next_observation=next_observation,
+                    terminated=terminated,
+                    truncated=truncated,
+                )
+                self.agent.add_experience(experience)
+                observation = next_observation
+                self.env.render()
path: toyrl/sarsa.py
user: login: coderabbitai[bot]
body: _⚠️ Potential issue_

**Call `env.render()` only when a render mode was requested**

Calling `env.render()` after creating the environment with
`render_mode=None` raises an exception in Gymnasium ≥ 0.28. Protect the call:

```diff
-if self.env.render:
-    self.env.render()
+if self.env.render_mode is not None:  # type: ignore[attr-defined]
+    self.env.render()
```

(or omit rendering entirely during training).  



> Committable suggestion skipped: line range outside the PR's diff.

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-04-27T14:20:49Z
author_association: NONE
diff_hunk: @@ -0,0 +1,241 @@
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import wandb
+
+
+class PolicyNet(nn.Module):
+    def __init__(
+        self,
+        env_dim: int,
+        action_dim: int,
+        action_num: int,
+    ) -> None:
+        super().__init__()
+        self.env_dim = env_dim
+        self.action_num = action_num
+        self.input_dim = env_dim + action_dim
+        self.output_dim = 1
+
+        layers = [
+            nn.Linear(self.input_dim, 128),
+            nn.ReLU(),
+            nn.Linear(128, self.output_dim),
+        ]
+        self.model = nn.Sequential(*layers)
+        # self.train()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)  # type: ignore
+
+
+@dataclass
+class Experience:
+    terminated: bool
+    truncated: bool
+    observation: Any  # S
+    action: Any  # A
+    reward: float  # R
+    next_observation: Any = None  # S'
+    next_action: Any = None  # A'
+
+
+@dataclass
+class ReplayBuffer:
+    buffer: list[Experience] = field(default_factory=list)
+
+    def __len__(self) -> int:
+        return len(self.buffer)
+
+    def add_experience(self, experience: Experience) -> None:
+        self.buffer.append(experience)
+
+    def reset(self) -> None:
+        self.buffer = []
+
+    def sample(self, with_next_sa: bool = True) -> list[Experience]:
+        if with_next_sa is False:
+            return self.buffer
+        else:
+            res = []
+            for i in range(len(self.buffer) - 1):
+                experience = self.buffer[i]
+                next_experience = self.buffer[i + 1]
+                res.append(
+                    Experience(
+                        observation=experience.observation,
+                        action=experience.action,
+                        reward=experience.reward,
+                        next_observation=next_experience.observation,
+                        next_action=next_experience.action,
+                        terminated=next_experience.terminated,
+                        truncated=next_experience.truncated,
+                    )
+                )
+            return res
+
+    def total_reward(self) -> float:
+        return sum(experience.reward for experience in self.buffer)
+
+
+class Agent:
+    def __init__(self, policy_net: PolicyNet, optimizer: torch.optim.Optimizer) -> None:
+        self._policy_net = policy_net
+        self._optimizer = optimizer
+        self._replay_buffer = ReplayBuffer()
+        self._action_num = policy_net.action_num
+
+    def onpolicy_reset(self) -> None:
+        self._replay_buffer.reset()
+
+    def add_experience(self, experience: Experience) -> None:
+        self._replay_buffer.add_experience(experience)
+
+    def get_buffer_total_reward(self) -> float:
+        return self._replay_buffer.total_reward()
+
+    def act(self, observation: np.floating, epsilon: float) -> tuple[int, float | None]:
+        if np.random.rand() < epsilon:
+            action = np.random.randint(self._action_num)
+            return action, None
+        x = torch.from_numpy(observation.astype(np.float32))
+        max_q = torch.tensor(-np.inf)
+        with torch.no_grad():
+            best_action = 0
+            for action in range(self._action_num):
+                x_ = torch.cat((x, torch.tensor([action], dtype=torch.float32)))
+                q = self._policy_net(x_)
+                if q > max_q:
+                    max_q = q
+                    best_action = action
+        return best_action, q.item()
+
+    def policy_update(self, gamma: float) -> float:
+        experiences = self._replay_buffer.sample()
+
+        observations = torch.tensor([experience.observation for experience in experiences])
+        actions = torch.tensor([experience.action for experience in experiences])
+
+        next_observations = torch.tensor([experience.next_observation for experience in experiences])
+        next_actions = torch.tensor([experience.next_action for experience in experiences])
+        rewards = torch.tensor([experience.reward for experience in experiences]).unsqueeze(1)
+        terminated = torch.tensor(
+            [experience.terminated for experience in experiences],
+            dtype=torch.float32,
+        ).unsqueeze(1)
+
+        # q preds
+        x_tensor = torch.cat((observations, actions.unsqueeze(1)), dim=1)
+        q_preds = self._policy_net(x_tensor)
path: toyrl/sarsa.py
user: login: coderabbitai[bot]
body: _⚠️ Potential issue_

**`torch.cat` will fail due to dtype mismatch; cast actions to float**

`actions` is `int64` by default, whereas `observations` is `float32`.  
Concatenating tensors of different dtypes raises:

```diff
-actions = torch.tensor([experience.action for experience in experiences])
+actions = torch.tensor(
+    [experience.action for experience in experiences], dtype=torch.float32
+)
```

Adjust the same for `next_actions`.  
Without this fix, training crashes on the first episode.  



<!-- suggestion_start -->

<details>
<summary>📝 Committable suggestion</summary>

> ‼️ **IMPORTANT**
> Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

`````suggestion
        observations = torch.tensor([experience.observation for experience in experiences])
        actions = torch.tensor(
            [experience.action for experience in experiences], dtype=torch.float32
        )

        next_observations = torch.tensor([experience.next_observation for experience in experiences])
        next_actions = torch.tensor([experience.next_action for experience in experiences])
        rewards = torch.tensor([experience.reward for experience in experiences]).unsqueeze(1)
        terminated = torch.tensor(
            [experience.terminated for experience in experiences],
            dtype=torch.float32,
        ).unsqueeze(1)

        # q preds
        x_tensor = torch.cat((observations, actions.unsqueeze(1)), dim=1)
        q_preds = self._policy_net(x_tensor)
`````

</details>

<!-- suggestion_end -->

<details>
<summary>🧰 Tools</summary>

<details>
<summary>🪛 GitHub Check: codecov/patch</summary>

[warning] 121-122: toyrl/sarsa.py#L121-L122
Added lines #L121 - L122 were not covered by tests

---

[warning] 124-127: toyrl/sarsa.py#L124-L127
Added lines #L124 - L127 were not covered by tests

---

[warning] 133-134: toyrl/sarsa.py#L133-L134
Added lines #L133 - L134 were not covered by tests

</details>

</details>

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-04-27T14:20:49Z
author_association: NONE
diff_hunk: @@ -6,6 +6,7 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
+
 import wandb
path: toyrl/reinforce.py
user: login: coderabbitai[bot]
body: _⚠️ Potential issue_

**Fix import sorting according to the pipeline failure.**

The imports need to be organized according to the project's standards. The pipeline is reporting: "Import block is un-sorted or un-formatted. Organize imports."


```diff
-
+
import wandb
```


> Committable suggestion skipped: line range outside the PR's diff.

<details>
<summary>🧰 Tools</summary>

<details>
<summary>🪛 GitHub Actions: Test</summary>

[error] 1-10: Ruff: Import block is un-sorted or un-formatted. Organize imports.

</details>

</details>

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-04-28T06:50:10Z
author_association: NONE
diff_hunk: @@ -175,19 +178,20 @@
             solved = total_reward > self.solved_threshold
             self.agent.onpolicy_reset()
             print(f"Episode {epi}, loss: {loss}, total_reward: {total_reward}, solved: {solved}")
-            wandb.log(
-                {
-                    "episode": epi,
-                    "loss": loss,
-                    "total_reward": total_reward,
-                }
-            )
+            if self.config.train.log_wandb:
+                wandb.log(
+                    {
+                        "episode": epi,
+                        "loss": loss,
+                        "total_reward": total_reward,
+                    }
+                )
path: toyrl/reinforce.py
user: login: coderabbitai[bot]
body: _🛠️ Refactor suggestion_

**Update wandb logging to check for availability.**

The wandb logging should check if wandb is available before trying to use it.


```diff
-            if self.config.train.log_wandb:
+            if self.config.train.log_wandb and _WANDB_AVAILABLE:
                 wandb.log(
                     {
                         "episode": epi,
                         "loss": loss,
                         "total_reward": total_reward,
                     }
                 )
```


> Committable suggestion skipped: line range outside the PR's diff.

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-04-28T06:50:10Z
author_association: NONE
diff_hunk: @@ -139,14 +142,14 @@
         self.gamma = config.train.gamma
         self.with_baseline = config.train.with_baseline
         self.solved_threshold = config.env.solved_threshold
-
-        wandb.init(
-            # set the wandb project where this run will be logged
-            project="Reinforce",
-            name=f"[{config.env.env_name}]lr={config.train.learning_rate}, baseline={config.train.with_baseline}",
-            # track hyperparameters and run metadata
-            config=asdict(config),
-        )
+        if config.train.log_wandb:
+            wandb.init(
+                # set the wandb project where this run will be logged
+                project="Reinforce",
+                name=f"[{config.env.env_name}]lr={config.train.learning_rate}, baseline={config.train.with_baseline}",
+                # track hyperparameters and run metadata
+                config=asdict(config),
+            )
path: toyrl/reinforce.py
user: login: coderabbitai[bot]
body: _🛠️ Refactor suggestion_

**Update wandb initialization to check for availability.**

The wandb initialization should check if wandb is available before trying to use it.


```diff
-        if config.train.log_wandb:
+        if config.train.log_wandb and _WANDB_AVAILABLE:
             wandb.init(
                 # set the wandb project where this run will be logged
                 project="Reinforce",
                 name=f"[{config.env.env_name}]lr={config.train.learning_rate}, baseline={config.train.with_baseline}",
                 # track hyperparameters and run metadata
                 config=asdict(config),
             )
```


> Committable suggestion skipped: line range outside the PR's diff.

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-04-28T06:50:10Z
author_association: NONE
diff_hunk: @@ -0,0 +1,244 @@
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+import wandb
path: toyrl/sarsa.py
user: login: coderabbitai[bot]
body: _⚠️ Potential issue_

**Fix import sorting according to the pipeline failure.**

The imports need to be organized according to the project's standards. The pipeline is reporting: "Import block is un-sorted or un-formatted. Organize imports."


```diff
from dataclasses import asdict, dataclass, field
from typing import Any

import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

import wandb
```


> Committable suggestion skipped: line range outside the PR's diff.

<details>
<summary>🧰 Tools</summary>

<details>
<summary>🪛 GitHub Actions: Test</summary>

[error] 1-10: Ruff: Import block is un-sorted or un-formatted. Organize imports.

</details>

</details>

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-04-28T06:50:10Z
author_association: NONE

PR comments:
user: login: coderabbitai[bot]
created_at: 2025-04-27T14:16:56Z
body: <!-- This is an auto-generated comment: summarize by coderabbit.ai -->
<!-- This is an auto-generated comment: rate limited by coderabbit.ai -->

> [!WARNING]
> ## Rate limit exceeded
> 
> @shenxiangzhuang has exceeded the limit for the number of commits or files that can be reviewed per hour. Please wait **1 minutes and 12 seconds** before requesting another review.
> 
> <details>
> <summary>⌛ How to resolve this issue?</summary>
> 
> After the wait time has elapsed, a review can be triggered using the `@coderabbitai review` command as a PR comment. Alternatively, push new commits to this PR.
> 
> We recommend that you space out your commits to avoid hitting the rate limit.
> 
> </details>
> 
> 
> <details>
> <summary>🚦 How do rate limits work?</summary>
> 
> CodeRabbit enforces hourly rate limits for each developer per organization.
> 
> Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.
> 
> Please see our [FAQ](https://docs.coderabbit.ai/faq) for further information.
> 
> </details>
> 
> <details>
> <summary>📥 Commits</summary>
> 
> Reviewing files that changed from the base of the PR and between 4ead5f11e42f1fcd132a3c89920b29ba778c6236 and a475952082887f6aa3319062d74f3fcc631baef1.
> 
> </details>
> 
> <details>
> <summary>📒 Files selected for processing (4)</summary>
> 
> * `.github/workflows/test.yaml` (1 hunks)
> * `.pre-commit-config.yaml` (1 hunks)
> * `toyrl/reinforce.py` (3 hunks)
> * `toyrl/sarsa.py` (1 hunks)
> 
> </details>

<!-- end of auto-generated comment: rate limited by coderabbit.ai -->
<!-- walkthrough_start -->

"""
## Walkthrough

A new file, `toyrl/sarsa.py`, introduces a complete SARSA reinforcement learning agent implementation using PyTorch and Gymnasium. The code defines a neural network policy, a replay buffer for experience storage, and an agent class that manages action selection and learning using the SARSA update rule. Configuration dataclasses are provided for environment and training settings. The `SarsaTrainer` class orchestrates the training loop, handles experience collection, applies policy updates, manages epsilon decay, and logs key metrics. The implementation is structured for extensibility and includes a main guard for running the training process.

## Changes

| File(s)                            | Change Summary                                                                                  |
|----------------------------------|------------------------------------------------------------------------------------------------|
| toyrl/sarsa.py                   | Added SARSA agent implementation: policy network, experience buffer, agent logic, configuration, trainer class, and training loop with logging and main guard. |
| toyrl/__init__.py                | Replaced module docstring with explicit imports and `__all__` export of SARSA and Reinforce classes. |
| toyrl/hello.py                  | Removed module containing simple hello and add functions.                                       |
| README.md                       | Added badges, installation instructions, updated documentation URL, and algorithm checklist including SARSA. |
| mkdocs.yml                     | Updated repository URL and name to ai-glimpse organization.                                     |
| pyproject.toml                 | Added mypy override to skip imports in wandb modules.                                           |
| .github/workflows/build_docs.yaml | Limited docs job OS matrix to ubuntu-latest only.                                               |
| tests/__init__.py              | Added empty test module docstring for toyrl tests.                                             |
| tests/test_sarsa.py            | Added comprehensive unit tests for SARSA components including policy net, replay buffer, agent, config, trainer, and minimal training run. |
| tests/test_reinforce.py        | Added comprehensive unit tests for Reinforce components and training.                          |
| tests/test_hello.py            | Removed simple test for removed hello module.                                                  |
| toyrl/reinforce.py             | Renamed Config to ReinforceConfig, added optional wandb logging control, and updated trainer accordingly. |

## Sequence Diagram(s)

```mermaid
sequenceDiagram
    participant Trainer as SarsaTrainer
    participant Env as Gymnasium Env
    participant Agent
    participant Buffer as ReplayBuffer
    participant Net as PolicyNet

    Trainer->>Env: reset()
    loop for each episode
        Env->>Agent: observation
        Agent->>Net: Q-values for all actions
        Agent->>Env: select action (epsilon-greedy)
        Env->>Env: step(action)
        Env->>Agent: next_observation, reward, terminated, truncated
        Agent->>Net: Q-values for next state-action
        Agent->>Buffer: add_experience()
        Agent->>Agent: policy_update()
        Agent->>Trainer: log metrics
        alt terminated or truncated
            Trainer->>Env: reset()
            Agent->>Buffer: reset()
        end
    end
```

## Assessment against linked issues

| Objective                                                         | Addressed | Explanation |
|-------------------------------------------------------------------|-----------|-------------|
| Implement SARSA RL agent with on-policy updates, experience buffer, epsilon-greedy, and training loop (#7) | ✅        |             |

## Poem

> In the meadow of code where the SARSA hops,  
> A rabbit trains agents with PyTorch props.  
> States and actions, rewards in tow,  
> Through CartPole fields, they learn and grow.  
> Epsilon dances, Q-values gleam,  
> On-policy dreams in a reinforced stream.  
> 🐇✨
"""

<!-- walkthrough_end -->
<!-- internal state start -->


<!-- DwQgtGAEAqAWCWBnSTIEMB26CuAXA9mAOYCmGJATmriQCaQDG+Ats2bgFyQAOFk+AIwBWJBrngA3EsgEBPRvlqU0AgfFwA6NPEgQAfACgjoCEYDEZyAAUASpETZWaCrKNxU3bABsvkCiQBHbGlcABpIcVwvOkgAIgAzEmouRGdU0NjIAHc0ZAcBZnUaejkI2BJIbERKe3KMAA94TCIAL1hsZvRkTx8/QODEXEgMRwEagE4ULFxyvu58RHV8Fw0YWaVEBgp4bnF8LH9E/wwGaUYvBfgMIhREBwqR5jG+AHZs9Vgy1FtVgDl8dC0WjqeD7NC+JS4bReZDLSAkercLyYaigjDIFT4PBlCoMWDNM5wxLUbD+ZBXXAURTYU70K44njeXz+IIhdD+HhUiTwJS0VZwCq2RiYSBjSDckhZGLUSCwXC4biIDgAemVRA+2AEGiYzGV2mIXngzEVJGVBFkFC8yp6VvGGiM+mM4CgZHo+HiOAIxDIymKClY7C4vH4wlE4ikMnkTCUVFU6i0OkdJig7nJGKwaDwhFI5Cofp1bAwnD8aCy9kczGc8lK0eUcc02l0YEMTtMBnNluVqQoqQ03FkHAMsWHBgskAAggBJb256gxBxOFz8D14gmINwIZDkMvxeDRFDG6KF3AYyAAZXHNgvfSu8WWp2PkGizgwVxuaBzQyqb+ssmg98+TB6AAcVkZgMFyeBHFWSchiUXdyFPAADKx8ENBhZF+EhcCQ4YSFJcE8NwLJlgAazKGUQiNOdkAARTACRwQGSB4ipZgFBOOcIL9QY53QDB6DQMQ0SmTwT3Cb9rnQFiSDoO8KByChBIoPF1DDUkKiyD5+HIWUeSULBkVkSh+VmJCAFFEUoeAyFOXDaGoITkTuewCDJKZgW5WgOl8SlMEWPZ0SmBgvGwYEpN4mhwiEwLwn8RTaHCch6iGSKSCShEhhitFooEiJKEKbiRPiZEiEQVZx0gJCbBIJE0FkAAhbB4iOXCmB8MNkARbhrNss4gOFLBUkPCoZhIditJmfg+Em2AsSGZLUqhGgwGy/YeG0HsWLhPyrjfe0DAFKrx0/NrnOQSsINIZAxp4ND4AwoiSIoUjwnwXYjXgFpKFy+h/Dq6tmqOGDUuwbh5goE94UVPd9mIfw6HkNahpIaJhPW0oSEY0LUSk+jseY+T0F6ZH03oPF8EuKTbrGQZpukqgBJYdB0YwUyKnmdD5DBhyaEqaobtmC8r0qqEKFIYt4ucegkSqSBgU2LEixifGmIqd0GUW1y51W1mNvgChovBrxZB/NgRUQIJnBiSgqT4C4XIGgQhNI3h3o/VF9gOgBhfZdyIAjAvlxyQtyAX7B6hh4HieQyG5KkMEfAbdtfKTuGcNA2BoHt2aqs80jQaAqCuSgzrDqYQXBL6zluuODf2Y9os/X6ImL1Obgud6Qb8bAgtqpBFGkcJ2rRqHut6k4h54Sh5OYbp7senmaPQeJs/hITPn7xBB/CJQGHqrqYYuDAW4uMrICz7YGHJE5QvCjuFkQcJzawVXQvS/j6AIKFmUlaX+Z/AAdRIPAIgcpkAADJIANSaALVY5luADw2PlCghU+LrTrgnR82c0FBx2hQXu+9YqfwZP9IyoogY1FQGSbCOkwCcwekjVeNQkh4mhkgkgB0jq1nLODZYUN9gm3lkgLY2EKjIwjkJfqeV4KZi8FDAgDJvbOFwKhaIDEACM8IMDxwbuwbuFISBEDzGcce2w+oKA6qzcIwtrxLxoE/aGixj7w1krQWOiILh5hyiQs+6opL0jQAocoiwpCt20O3HgOxUYlwOuYSw455G+jRDdAEt097Im8fsWEHpur8JiHCTwAh0LaMiDZdcBgoDjiBDEUOLkUIL0wthXC9IkIditN2Xs/ZcKzQ4oMAhYg4QDSQvJBKuEs5zT5JUicNT6C8ycuXCyVlzGTxaVgNp+ALQdILn2WQPTtK7lRrQJUKDcF0HCJSQhc5EohmqBQRixDkZxT/kpDKKUAD6gg7kPJ8Ytd5yN7RVNmcHKEdTkDVVqkZJqLVS5TCqu0rsOzunvCmhMxQJykLvPedEDAWKkLhCQmgIE7yzE2VWQSmhOECXDSRCQfF8L8A/3eVLJSSFAUzN5OcRZJ12BrIZVsxFPY0C7P2ai7CkyMVYr2rgPFBL9gMIwsy6QzSCVEtoCS5ZZK7IEolu8gQlCKDvO/uCJVYzVViHpUhBVsh3n2Lpey6pnL5lgqquZHRvsMD+z5RsgVnThXIt6YcrwxyuBx3eRBNgzyBKUHecwHe9g0JSHVTMMkc1g0OuBc686VUi4RI9V6uFPrOx+pFSiz4QaQ2QCIJnSsSVHAksQdvDY4RnwUHbsyucLb8BEHeTkASWppmOpiFmxZ+chX5tAd6hFJaA0HJssGk5ccLltwzZyl1SEx2pFzSXCgU7NnFqRXsstF9xXoq4Ji950rZXwrbmywdwLKz0gDv/XpCJM60v5h+XE/hcY3A3QXCdRBcIDQIRgSJ/6hXbtzBoFOAAKAAlHe4csQHRtg0OqGYmplTPVIiVfAWREDKn1XudVtB8DXw0LITOXhBzIdHAk6cOZfTzgrFWZcjB8TXGkBuTSZE8NlkrJSeA9Rtp8FukhMj19cJCEEHC26wF1AAAlNQTlZsgJgnrQGBxEjkC6iho42RKFGFgx4fzzQZEhSsTBEBgGRA4nCJCkJaSZgR2zNEHPvWSRFWQgxxqlPMeVNYqBVxcb6P0h6CjZgyYEBEAEoGdJCPWkhTUvdcDYDc/Z3CRNbqIEjgZmIVhZAzHWlIHsIkADMGgNGVcgP8fgY0+AheurFhk0XXK1UcRp4EgVHFwjPg9bIlAKiViUHE+jE4klZKCoo9JohMmeyChrPJkMCl8CKSU9gIJuMGH+OQe0dHKlthsOZccAARAAsuZDQzBaC0ZHGOKcM5mP0AXJWJcGsmvbaOsds7l3rv0F3PuXT2jOO0lFEjIEZtvDiA/c7WgzWZS3QINwcIVwQphR/IV4rWBSuLHWg4PhkNwiFasJODaDBSJfvFJQPHJ95b4YwBcIl2txCDAeo4hTuBlMxfQmQaow9B6QCtoaXA8hYM2CBvBi5sgeocdEKRH8sHmAy9kFLzkJAwAFnUCzqou9yOOHYAtihJGdeOPs9kPjFwyyRV1yQnhTBStftzpJg3RYjcAFUbAABlsi5EqNwXmMRWLMyCT1HsYJfCc+59YL9eQ1ItduvqIghpjTVEgFHlTVhY/2DUhVSHgVwQmwytbsMKTRNTF4j4I3Vx+k0h65AWDvTQ87Arz/XwBYgLwZIUEvE8vDR0w1kaWlx5pReCIMsD4c8G81VvPecReVbGVTe6RWpLBaVRUgKd2ivwL7OBX3SE4a/og0C71KDkaq6BjYe5NhbqSGQZIzvXpbiJ8lujW5qDbRYtsVN2yQIcI5DswAjBmBSJJNyowIaN/8UMHtGMfQTEXtWN3sVxONroeM5hLg3J5BPcfcQNaoMDlh5AI0Kh6RbpztSJTtyN1M/YtMpsWI9xNIht/dA8+Q1gOZ/BuQsRkACYzh5gKQfxFFYg5QFQlRVQMN2gtQdQuw6hGhmg2gOhrgzR90vBMgBpYhEBpCmhrg5DmhFCtlMgz9cQUCYgBChDFQVQ1QNQJCWA9R4ADQh9qhdDLQVC8pYgk8U8TRHDlCwtI5wxUZZA/gARGVyhGtqCA5aC4Qs4iVHJfcMRJtL8jB4kJts5b8E91g5tH8y9n8IY/RCkP8BtNtxBttf99sACIAgCDB+w3YRAxAYMWBIC6MYCnt4DyxFx5APsjCKkjokItdxBvoMBcJ04Kcqdgd/A40k0WI2ILMlAJBcJ3oC8IRapXRbJyknwkBNBWDOR8BqihhsDkAiZFMWBaov1wgmobhc0KcfpN99djwFsW4zwsQVIKhfYlBBsOQ7UAdJjxN1CyAZCtD2hmhcIM8YsBogSY9mtdjUiKhCVbDk97C6V08lMVMQTESYss9wTvdAsuiABtJQHqaNE4WQeGLERUAAXVwmqD1mB0WEPHyyMz6Xul5n4LMmmNwmMWJKmEUSCUWGuHUUNF0n73WKOk2w5CQixIIDQm1BMxBD6LJI4n9m0wxguApxiL6DGMvwnDwgExV34FK22BeOBwvy/gBEQAV24G2h8Hwx/CH34T2KGV6DjW8miAumoFUikicyAi1AACo2UNikJQEMBlh4SBTKho0+BRTxSvANACEWpZTgdPsPiQ8L5odbC+SKhgzOSc8eSNdUzVi6YNMoQ9o3TYhTtYg70Ejr9kj68Zs0jQ4pscl4QX8Vs39GRil8iv9Cif99hOEoDUNyj7NCMpVXwZV3ldk7toCGMmjrkWi3s2jkC1wjBKptw6D9w2kQgBzL0hysVS18QZBZIMxgV8yIkfwuS3x9wcywChNqZ8QhgeRNtDlBZhtqR9w/cglzcHTvAKgiZYh2kNBMg6t4hCF5iFAXi+tu0BtUAKQqRvJaQr8Ekb8qy0kaz5sn9clGycj39WyGBSlv8jBf8ezACjB+yzQQh3lygLTRyezGimNmjXs2N2j5zDo0jj8g96COICzyBBJMzk9Ro2QAKTgg5VzBhSLUYLhcIZgZRrSmyLMyLRKWJALipPizIfyZL8BxknzOENjzcERRA8BjCzIVKst5L9hwhBhAy6QoZ/A0s21xQ1Y5NZhGJtgVAVy54gMW4w5KA/RxKhhbpLLSQOKbL34VSRQa8oRJ52MctLybhRcepc5zcqTdowFcBRlpZwhSto5TYpJBF5AkcZd1YPRbp5oxIy1zM1UQQI85cKdYQ+BbY4RONaA+SiBYKkjklslIT5Z0i6z2Nlt0KWzP8ykiiuySiUMCL2w1ziKhL/BZ9HiKKGiJzqKpzaKkCOMGLUxNT8o6Z3z9w3YvIzgdReASA6hQkKhe5td+zy8fKQFPU59iZx9tgZh2Ih8jxDdAoDFb4woa41zy8mAOQ9quyiw8gaRAJwU1FGEsIqUqoaoAZoVWoCUeUix6UNN5TaCwUzhYMIUprThANLU3UJAsaCUoMsa1dk4zIZ8rrHioNS5YrPq0qY5RRIIsK+LWYq5RdViV9RMEoNo7gzgSDZhrUnoyJwhSULFTKqBSASEWVBIGAGBHBvBq9phZgyF6oKEYVDZ0BPxGAf0g4BoJEaUGqW5EaaCjd04qAs4ahZFocyYdTbZbzMTZgU4qEboQgg9CkqQw8K5xAq4WgjdtaL49pKxfI24fx4sNYuz2Em1RoAQ+cNIGQU4fxO4zTtLpaHFirsQaqc5NLPqqhxEubIYy9FEDbUEGyfDjDcrHF1C0AepHEBpuCq68oHcagBAgj2r4g5Eso66qgCB2IDawjb9c58BHdeh8yqRfB+M5NPqrgJBE0zhqg0szSQ7QwxBHEJ78AFcpI0VjkW5sYeRf1+A8AdQp4o7tgryKhJqybTgbqJ97r/R5hyB/rRQDq0Awk/dDFo16Be4Ywc8Rp7BTgIJtgFgmrElKy87EKKgH9OqsjX96Z1s2z+qKkoAsIyxzdGb69DTC0iL7MlUMbOFulBxIBdB4USLrVw1sIEMkIDA8GoBBKZVFabV9UVbSHyH8GqH/lPx3lREFsGGKGCGhKv0ix/kxBOGmGMHu7BHKGMH7bDV2HApRHuGZVCpqIvAjVA7rgGGDsyjCKxqMGZ0BxKK5q4CFrEDZzlquNOjNw1rzcHB49sshZLxrxwRbrJ8FBcckz5Edh9xfrb6BF8rZhHrxpnq0RXr0dkEzrrGoSQaMIwbcJyACJfByBiIBbqdtgY4mTPzlgOb04XJCrsRy7Zdk5crc4IUoaDUWlHbfMAc4Rvr/AxB778QOCZo6h4RjRRdopmE+BSqpIhbJ4y730Grj0tY0pdYg504DZa6/plV5R9pfS4aHMXVUB7Nna1tXaqEhymhDQvata8pxMYpxlT1aAAByXc2p0EPgd+moYxNxMXDrWGDAAAXgAAYu8QMgJmZYMrnj4bmNEu8FVylBb0RSQfwt7Ja8E8AxJLa35gg+grL0QCnN00AsauUXI5mnbym+Bzb5EArmIBppbTKu7QiFSsB+7raNgYWC4Kbd1bh1q/QiZKmwxW9MAPag5X0dF65E52Bu9m5u9fbXx/bwlCybh4stL6gdLk6Tysyw7BdFEabsq7blGH53oe4gpZpzM07baeLBhkAs708wIIJFhHBtFdFWW76BpCt/wVJPgaB0RlhRmcQoxyhlToqp4cmp5qAhN9U+Z9rqhJ4W5x4xAYga6WtAW+JbpF9r6/qTw9mFAKAqncBEJLbDFjEFsAH4LgH78OqUiIGpLcjMLsKOyHRatJRKW5L+Ky9UHWl0GSLtGkJcGhHCHGliHcAZHmGaG9UDVG2MHeGZUpG0Q22SKO3+GG3ENGGxGSKRHB2uHmGJG2HNbu2x2a2hL5H/alGjyVHB21G2wEVBz1Atz+wxzxtHt5qeJDH2NPtTHHzHTQHKDIqqpYhzsrArBoBIBzt1Lfyek/dRjCXBI8pFawdMZPEHptdJKvGIbLr5JMbQjLVSbQOSAyWstFKoTT6oPyTNRNr596AINUg8a85SW25YVg92IviC4kOChn2OUyqIIfBZBcoqosVC88VcyhhgdILqQwdFFup0J1AhEwHRoQlUnSQDxsiYgiCMQHzeqBtxxSdwqlKlCBiXYncEi93k3WrqzQG02UKi7IGs2+qcL71OU2P/2hhuquCmhqP/kfA6Oy2lDlRN3hzS0BS3xz1IO58YOCVHPHjMP0PC4cPd0CVYWsahrezCLLOVKZr7s9HZxD3Wjj2Oi0DhSoSEVgvkUUOVT32k1TJUAkvDyS4v4SJP6P1kGUlz0VLSHHM1VYM0BwgBBEMCmivEMi29ZfK209KKhwspJYhFMRL8BwhAFlhg0QBnCD8gnpR6dr4r2m92rNhtgPpjLIW/KIhcqW4qgqc30RpVWqpSvyvRRav8v1pxTyJ4mWImdxBrh6F8AKRhhRgacuh2QA5jxLaGv/KxoDYWjooYQAQXcR96Am9xZXc4JpAtgdgetnkoW5ueoFvUgxblvaVAsGom7tvFXGD+ltBErkrXk+mAQ9rogRMh60IDv8N6Y07ZQgIGqk2gGlOQH2raz03UKBPmzoGsKCjyk82ao1SAcjL1kavdA9BXJyW8P+VOwEu9lplmeP26uBL1uuA8NqAKuJfDuu9mxcfqBYPmYi0rQBeyG12+zLOEO58QvxyJxYDwuWNIv6KTG0Du78WEXwV4WRiyBM5jCAR0az6SA/ONSlzG60IkgsAK0qoz5e0PTX2MRgVFEkICbwOWtsffAshygGtIBgFQFwFIAoEYFchUa+1aBKunxu1/EbhUBbfil1SujXPTgYO+lLlBk+AAzrc3Ia5eb3LyZQjC1qgvB4hJTNMgM7d9hus0RC95BpVVnq5kA0+YssqUAPQkJfeh+SnW5ghVgzxPpMki8GQh/M+iBs/hRega9byY7ZXM/5XUBK+OIu/yqksJ//fc5H0hp/vdg6Wv9l5HeoP4WQqkg3Qx/4WBoZ7wVT/+0bmi5gg+Ukcd+3ddmGnhPaAg3QU3cjr7n7QckjEyNfYJSDQjRA6St0buinSyiTYUmkbUaDv36wMASeLVabOTzAZU91OmbDClp1zbTJvY2aY4Hb1uxVR4WgAJMJgOWDeFrBgs4Cpte01bpPBh04sUjkPvbtH737RcB3evgG5pAAABi4IaoGJQd6h82+9kEONmnYHrIEUXA04CKl4FQAPUtecvjniIDcRo67xc9PBEgDWcsUsGJvvEAFxt8uAgGOXlz1/y4RmBEmEgB6AsHvIrBqMGwXKVARcAi+zvUIo4Nqxdk+U66QISX1UF88rQGg7Bnsl4Ea8jAfYfwJrilK4A0hbfSjNRl3ZUV9GEXGclFxWq81UhPRPwT3SDihNIAEuFqKsWVhtM8os8F1jUGBzvFwcxjNfrdFxwiRDgQ2MKrzyQgSA7mVWDRBoAABscgqqEMJGEaAXgPpOrEERqAOkDMRCFNnNBXpV1vut3N6HwH+j4EAsrxYbESm7Llk4KpPQgam0p5qduqq2UTvT3bKM9pkdWUAaxzQq3C6eObcpBG2IE9Z/Oh2FsMmG0Qv9PQ2YfIavgDBFguAVAa3EexrCDxYwagBsImGbAGBARPRS9MciVQSgpQ6qXiJDCbAAjnQkAMYXczubjA7maANAOMFoAAAWBgLQDGE0iNEtABgIyJagAAmBgC3XZG0B4gLwDRGgBeAABWIURogFGki3gSYIkTSLoAaIaRQolqDSJJECAZRAADjJHxAVRLwNAKqKJQ0i0A8QRkeMCFECANEQo9wS8GFEEjURRIsYWMNVG0jxg8QDRGMPJGqj2RAgBgNqPZHlZysNI8rPEDuZCi7mqo50WMJIBOjnRDAGkeMFdHWjARJAYMQaNoCqimRdzAQEoHTFoBysAgeIDSNTHxBgxQohgGSLGEMBys4wcrEKM5ECBVRQo+MdKOf4KixRJAGkeyOdHxB6R1WdkdmIYCqjxg4wdkemPZHjBnYlo1UayJ9FjDrRNoiAP6EKAyoeQiALETZBxEko8o+gIAA -->

<!-- internal state end -->
<!-- finishing_touch_checkbox_start -->

<details open="true">
<summary>✨ Finishing Touches</summary>

- [ ] <!-- {"checkboxId": "7962f53c-55bc-4827-bfbf-6a18da830691"} --> 📝 Generate Docstrings

</details>

<!-- finishing_touch_checkbox_end -->
<!-- tips_start -->

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

<details>
<summary>❤️ Share</summary>

- [X](https://twitter.com/intent/tweet?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A&url=https%3A//coderabbit.ai)
- [Mastodon](https://mastodon.social/share?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A%20https%3A%2F%2Fcoderabbit.ai)
- [Reddit](https://www.reddit.com/submit?title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&text=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code.%20Check%20it%20out%3A%20https%3A//coderabbit.ai)
- [LinkedIn](https://www.linkedin.com/sharing/share-offsite/?url=https%3A%2F%2Fcoderabbit.ai&mini=true&title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&summary=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code)

</details>

<details>
<summary>🪧 Tips</summary>

### Chat

There are 3 ways to chat with [CodeRabbit](https://coderabbit.ai?utm_source=oss&utm_medium=github&utm_campaign=ai-glimpse/toyrl&utm_content=9):

- Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  - `I pushed a fix in commit <commit_id>, please review it.`
  - `Generate unit testing code for this file.`
  - `Open a follow-up GitHub issue for this discussion.`
- Files and specific lines of code (under the "Files changed" tab): Tag `@coderabbitai` in a new review comment at the desired location with your query. Examples:
  - `@coderabbitai generate unit testing code for this file.`
  -	`@coderabbitai modularize this function.`
- PR comments: Tag `@coderabbitai` in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  - `@coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.`
  - `@coderabbitai read src/utils.ts and generate unit testing code.`
  - `@coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.`
  - `@coderabbitai help me debug CodeRabbit configuration file.`

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

### CodeRabbit Commands (Invoked using PR comments)

- `@coderabbitai pause` to pause the reviews on a PR.
- `@coderabbitai resume` to resume the paused reviews.
- `@coderabbitai review` to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
- `@coderabbitai full review` to do a full review from scratch and review all the files again.
- `@coderabbitai summary` to regenerate the summary of the PR.
- `@coderabbitai generate docstrings` to [generate docstrings](https://docs.coderabbit.ai/finishing-touches/docstrings) for this PR.
- `@coderabbitai generate sequence diagram` to generate a sequence diagram of the changes in this PR.
- `@coderabbitai resolve` resolve all the CodeRabbit review comments.
- `@coderabbitai configuration` to show the current CodeRabbit configuration for the repository.
- `@coderabbitai help` to get help.

### Other keywords and placeholders

- Add `@coderabbitai ignore` anywhere in the PR description to prevent this PR from being reviewed.
- Add `@coderabbitai summary` to generate the high-level summary at a specific location in the PR description.
- Add `@coderabbitai` anywhere in the PR title to generate the title automatically.

### CodeRabbit Configuration File (`.coderabbit.yaml`)

- You can programmatically configure CodeRabbit by adding a `.coderabbit.yaml` file to the root of your repository.
- Please see the [configuration documentation](https://docs.coderabbit.ai/guides/configure-coderabbit) for more information.
- If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: `# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json`

### Documentation and Community

- Visit our [Documentation](https://docs.coderabbit.ai) for detailed information on how to use CodeRabbit.
- Join our [Discord Community](http://discord.gg/coderabbit) to get help, request features, and share feedback.
- Follow us on [X/Twitter](https://twitter.com/coderabbitai) for updates and announcements.

</details>

<!-- tips_end -->
author_association: NONE
user: login: codecov[bot]
created_at: 2025-04-27T14:18:55Z
body: ## [Codecov](https://app.codecov.io/gh/ai-glimpse/toyrl/pull/9?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ai-glimpse) Report
Attention: Patch coverage is `92.77108%` with `12 lines` in your changes missing coverage. Please review.
> Project coverage is 95.07%. Comparing base [(`b927268`)](https://app.codecov.io/gh/ai-glimpse/toyrl/commit/b927268af88992288e417dd90826aeda2c63cb63?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ai-glimpse) to head [(`a475952`)](https://app.codecov.io/gh/ai-glimpse/toyrl/commit/a475952082887f6aa3319062d74f3fcc631baef1?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ai-glimpse).
> Report is 1 commits behind head on master.

| [Files with missing lines](https://app.codecov.io/gh/ai-glimpse/toyrl/pull/9?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ai-glimpse) | Patch % | Lines |
|---|---|---|
| [toyrl/sarsa.py](https://app.codecov.io/gh/ai-glimpse/toyrl/pull/9?src=pr&el=tree&filepath=toyrl%2Fsarsa.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ai-glimpse#diff-dG95cmwvc2Fyc2EucHk=) | 94.15% | [9 Missing :warning: ](https://app.codecov.io/gh/ai-glimpse/toyrl/pull/9?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ai-glimpse) |
| [toyrl/reinforce.py](https://app.codecov.io/gh/ai-glimpse/toyrl/pull/9?src=pr&el=tree&filepath=toyrl%2Freinforce.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ai-glimpse#diff-dG95cmwvcmVpbmZvcmNlLnB5) | 66.66% | [3 Missing :warning: ](https://app.codecov.io/gh/ai-glimpse/toyrl/pull/9?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ai-glimpse) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           master       #9       +/-   ##
===========================================
+ Coverage   45.66%   95.07%   +49.40%     
===========================================
  Files           2        3        +1     
  Lines         127      284      +157     
===========================================
+ Hits           58      270      +212     
+ Misses         69       14       -55     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/ai-glimpse/toyrl/pull/9?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ai-glimpse).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ai-glimpse).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE