Repository: ai-glimpse/toyrl
Pull request #22: feat: add A2C
Author: login: shenxiangzhuang
Created: 2025-05-03T06:12:45Z | Merged: 2025-05-03T08:38:46Z

Close #23 

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

- **New Features**
  - Introduced Advantage Actor-Critic (A2C) reinforcement learning training with configurable environment and training parameters.
  - Added on-policy experience collection, Generalized Advantage Estimation (GAE), and entropy-based exploration to improve learning.
  - Integrated detailed training progress logging and metrics visualization support.
  - Provided a command-line interface for easy execution of A2C training sessions.
- **Tests**
  - Added comprehensive unit and integration tests covering neural network, experience buffer, agent behavior, configuration, and training workflow.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Linked issue #23: A2C

Diff:
diff --git a/tests/test_a2c.py b/tests/test_a2c.py
new file mode 100644
index 0000000..d1d20bd
--- /dev/null
+++ b/tests/test_a2c.py
@@ -0,0 +1,257 @@
+"""Tests for the A2C algorithm."""
+
+import gymnasium as gym
+import numpy as np
+import torch
+
+from toyrl.a2c import (
+    A2CConfig,
+    A2CTrainer,
+    ActorCriticNet,
+    Agent,
+    EnvConfig,
+    Experience,
+    ReplayBuffer,
+    TrainConfig,
+)
+
+
+def test_actor_critic_net():
+    """Test the ActorCriticNet class."""
+    env_dim, action_num = 4, 2
+    net = ActorCriticNet(env_dim=env_dim, action_num=action_num)
+
+    # Test forward pass
+    x = torch.randn(env_dim, dtype=torch.float32)
+    policy_logits, value = net(x)
+
+    assert policy_logits.shape == torch.Size([action_num])
+    assert value.shape == torch.Size([1])
+    assert isinstance(policy_logits, torch.Tensor)
+    assert isinstance(value, torch.Tensor)
+    assert policy_logits.dtype == torch.float32
+    assert value.dtype == torch.float32
+
+
+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], dtype=np.float32),
+        action=0,
+        reward=1.0,
+        next_observation=np.array([0.1, 0.1, 0.1, 0.1], dtype=np.float32),
+        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], dtype=np.float32),
+        action=1,
+        reward=0.5,
+        next_observation=np.array([0.2, 0.2, 0.2, 0.2], dtype=np.float32),
+        terminated=False,
+        truncated=False,
+    )
+    buffer.add_experience(exp2)
+
+    assert len(buffer) == 2
+    assert buffer.total_reward() == 1.5
+
+    # Test sampling
+    samples = buffer.sample()
+    assert len(samples) == 2
+    assert samples[0].action == 0
+    assert samples[1].reward == 0.5
+
+    # Test reset
+    buffer.reset()
+    assert len(buffer) == 0
+    assert buffer.total_reward() == 0
+
+
+def test_agent_creation():
+    """Test creating an agent."""
+    env_dim, action_num = 4, 2
+    net = ActorCriticNet(env_dim=env_dim, action_num=action_num)
+    optimizer = torch.optim.Adam(net.parameters(), lr=0.01)
+
+    agent = Agent(net=net, optimizer=optimizer)
+
+    # Check agent properties
+    assert hasattr(agent, "_net")
+    assert hasattr(agent, "_optimizer")
+    assert hasattr(agent, "_replay_buffer")
+
+
+def test_agent_act():
+    """Test the agent's act method."""
+    env_dim, action_num = 4, 2
+    net = ActorCriticNet(env_dim=env_dim, action_num=action_num)
+    optimizer = torch.optim.Adam(net.parameters(), lr=0.01)
+
+    agent = Agent(net=net, optimizer=optimizer)
+
+    # Test act method in training mode
+    observation = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32)
+    action = agent.act(observation)
+
+    assert isinstance(action, int)
+    assert action in range(action_num)  # For CartPole, should be 0 or 1
+
+    # Test act method in eval mode
+    action = agent.act(observation, eval=True)
+    assert isinstance(action, int)
+    assert action in range(action_num)
+
+
+def test_config():
+    """Test the Config class."""
+    config = A2CConfig()
+
+    # Check default values
+    assert config.env.env_name == "CartPole-v1"
+    assert config.train.gamma == 0.999
+    assert config.train.lambda_ == 0.98
+    assert config.train.value_loss_coef == 0.5
+    assert config.train.policy_loss_coef == 0.5
+    assert config.train.entropy_coef == 0.01
+    assert config.train.eval_episodes == 10
+    assert config.train.eval_interval == 100
+
+    # Test custom config
+    custom_config = A2CConfig(
+        env=EnvConfig(env_name="MountainCar-v0", solved_threshold=90.0),
+        train=TrainConfig(
+            num_episodes=1000,
+            learning_rate=0.005,
+            gamma=0.99,
+            lambda_=0.95,
+            value_loss_coef=0.7,
+            policy_loss_coef=0.8,
+            entropy_coef=0.02,
+            eval_episodes=5,
+            eval_interval=50,
+        ),
+    )
+
+    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
+    assert custom_config.train.gamma == 0.99
+    assert custom_config.train.lambda_ == 0.95
+    assert custom_config.train.value_loss_coef == 0.7
+    assert custom_config.train.policy_loss_coef == 0.8
+    assert custom_config.train.entropy_coef == 0.02
+    assert custom_config.train.eval_episodes == 5
+    assert custom_config.train.eval_interval == 50
+
+
+def test_trainer_creation():
+    """Test creating a trainer."""
+    config = A2CConfig(
+        env=EnvConfig(env_name="CartPole-v1", render_mode=None),
+        train=TrainConfig(num_episodes=10, learning_rate=0.01, log_wandb=False),
+    )
+
+    trainer = A2CTrainer(config)
+
+    assert isinstance(trainer.env, gym.Env)
+    assert isinstance(trainer.agent, Agent)
+    assert hasattr(trainer, "num_episodes")
+    assert trainer.num_episodes == 10
+    assert trainer.gamma == 0.999
+    assert trainer.lambda_ == 0.98
+    assert trainer.value_loss_coef == 0.5
+    assert trainer.policy_loss_coef == 0.5
+    assert trainer.entropy_coef == 0.01
+
+
+def test_agent_net_update():
+    """Test the agent's net_update method."""
+    env_dim, action_num = 4, 2
+    net = ActorCriticNet(env_dim=env_dim, action_num=action_num)
+    optimizer = torch.optim.Adam(net.parameters(), lr=0.01)
+
+    agent = Agent(net=net, optimizer=optimizer)
+
+    # Add some experiences
+    exp1 = Experience(
+        observation=np.array([0.0, 0.0, 0.0, 0.0], dtype=np.float32),
+        action=0,
+        reward=1.0,
+        next_observation=np.array([0.1, 0.1, 0.1, 0.1], dtype=np.float32),
+        terminated=False,
+        truncated=False,
+    )
+    agent.add_experience(exp1)
+
+    exp2 = Experience(
+        observation=np.array([0.1, 0.1, 0.1, 0.1], dtype=np.float32),
+        action=1,
+        reward=0.5,
+        next_observation=np.array([0.2, 0.2, 0.2, 0.2], dtype=np.float32),
+        terminated=True,
+        truncated=False,
+    )
+    agent.add_experience(exp2)
+
+    # Test net_update
+    loss, entropy, adv_mean = agent.net_update(
+        gamma=0.99, lambda_=0.95, value_loss_coef=0.5, policy_loss_coef=0.5, entropy_coef=0.01
+    )
+
+    assert isinstance(loss, float)
+    assert isinstance(entropy, float)
+    assert isinstance(adv_mean, float)
+
+    # Check that buffer is still intact
+    assert len(agent._replay_buffer) == 2
+
+
+def test_minimal_training():
+    """Test minimal training run with a single episode."""
+    # Create minimal config with just one episode
+    config = A2CConfig(
+        env=EnvConfig(env_name="CartPole-v1", render_mode=None),
+        train=TrainConfig(num_episodes=1, learning_rate=0.01, log_wandb=False),
+    )
+
+    # Initialize trainer
+    trainer = A2CTrainer(config)
+
+    # Run training for one episode
+    trainer.train()
+
+    # If we got here without errors, test passed
+    assert True
+
+
+def test_trainer_evaluation():
+    """Test the evaluation functionality of the trainer."""
+    config = A2CConfig(
+        env=EnvConfig(env_name="CartPole-v1", render_mode=None),
+        train=TrainConfig(num_episodes=1, learning_rate=0.01, log_wandb=False, eval_episodes=2, eval_interval=1),
+    )
+
+    trainer = A2CTrainer(config)
+
+    # Test evaluation method
+    eval_reward = trainer.evaluate(num_episodes=2)
+
+    assert isinstance(eval_reward, float)
+    # For a randomly initialized agent, we expect some reward but usually less than 50
+    # per episode for CartPole
+    assert 0 <= eval_reward <= 100
diff --git a/toyrl/__init__.py b/toyrl/__init__.py
index b0930ed..b3a56ab 100644
--- a/toyrl/__init__.py
+++ b/toyrl/__init__.py
@@ -1,8 +1,11 @@
 from .dqn import DqnConfig, DqnTrainer
 from .reinforce import ReinforceConfig, ReinforceTrainer
 from .sarsa import SarsaConfig, SarsaTrainer
+from .a2c import A2CConfig, A2CTrainer
 
 __all__ = [
+    "A2CConfig",
+    "A2CTrainer",
     "DqnTrainer",
     "DqnConfig",
     "ReinforceTrainer",
diff --git a/toyrl/a2c.py b/toyrl/a2c.py
new file mode 100644
index 0000000..1df5412
--- /dev/null
+++ b/toyrl/a2c.py
@@ -0,0 +1,264 @@
+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 ActorCriticNet(nn.Module):
+    def __init__(self, env_dim: int, action_num: int) -> None:
+        super().__init__()
+        self.env_dim = env_dim
+        self.action_num = action_num
+        self.shared_layers = nn.Sequential(
+            nn.Linear(env_dim, 64),
+            nn.ReLU(),
+        )
+        self.policy_head = nn.Linear(64, action_num)
+        self.value_head = nn.Linear(64, 1)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        x = self.shared_layers(x)
+        policy_action_logits = self.policy_head(x)
+        v_value = self.value_head(x)
+        return policy_action_logits, v_value
+
+
+@dataclass
+class Experience:
+    observation: Any
+    action: Any
+    reward: float
+    next_observation: Any
+    terminated: bool
+    truncated: bool
+
+
+@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) -> list[Experience]:
+        return self.buffer
+
+    def total_reward(self) -> float:
+        return sum(experience.reward for experience in self.buffer)
+
+
+class Agent:
+    def __init__(self, net: nn.Module, optimizer: torch.optim.Optimizer) -> None:
+        self._net = net
+        self._optimizer = optimizer
+        self._replay_buffer = ReplayBuffer()
+
+    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.ndarray, eval: bool = False) -> int:
+        x = torch.from_numpy(observation.astype(np.float32))
+        with torch.no_grad():
+            action_logits, _ = self._net(x)
+        next_action_dist = torch.distributions.Categorical(logits=action_logits)
+        action = next_action_dist.sample()
+        if eval:
+            action = next_action_dist.probs.argmax(dim=-1)
+        return action.item()
+
+    def net_update(
+        self, gamma: float, lambda_: float, value_loss_coef: float, policy_loss_coef: float, entropy_coef: float
+    ) -> tuple[float, float, float]:
+        experiences = self._replay_buffer.sample()
+
+        observations = torch.tensor([exp.observation for exp in experiences])
+        actions = torch.tensor([exp.action for exp in experiences])
+        rewards = torch.tensor([exp.reward for exp in experiences]).unsqueeze(1)
+        next_observations = torch.tensor([exp.next_observation for exp in experiences])
+        terminateds = torch.tensor([exp.terminated for exp in experiences], dtype=torch.float32).unsqueeze(1)
+
+        # calculate predicted V-values
+        policy_action_logits, v_values = self._net(observations)
+        # n-step return
+        v_targets = torch.zeros_like(rewards)
+        for t in reversed(range(len(experiences) - 1)):
+            v_targets[t] = rewards[t] + gamma * v_targets[t + 1] * (1 - terminateds[t])
+        # calculate value loss
+        value_loss = nn.functional.mse_loss(v_values, v_targets)
+
+        # calculate advantages by GAE
+        with torch.no_grad():
+            _, v_values_next = self._net(next_observations)
+        deltas = rewards + gamma * v_values_next * (1 - terminateds) - v_values
+        advantages = deltas.clone()
+        for t in reversed(range(len(experiences) - 1)):
+            advantages[t] = deltas[t] + gamma * lambda_ * advantages[t + 1] * (1 - terminateds[t])
+        advantages = advantages / (advantages.std() + 1e-8)
+        advantages = advantages.detach()
+
+        action_dist = torch.distributions.Categorical(logits=policy_action_logits)
+        action_entropy = action_dist.entropy().mean()
+        action_log_probs = action_dist.log_prob(actions)
+        # calculate policy loss
+        policy_loss = -action_log_probs * advantages
+        policy_loss = torch.mean(policy_loss)
+
+        loss = value_loss * value_loss_coef + policy_loss * policy_loss_coef - entropy_coef * action_entropy
+
+        # update
+        self._optimizer.zero_grad()
+        loss.backward()
+        self._optimizer.step()
+        return loss.item(), action_entropy.item(), advantages.mean().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
+    lambda_: float = 0.98
+    value_loss_coef: float = 0.5
+    policy_loss_coef: float = 0.5
+    entropy_coef: float = 0.01
+
+    num_episodes: int = 500
+    learning_rate: float = 0.01
+
+    eval_episodes: int = 10
+    eval_interval: int = 100
+    log_wandb: bool = False
+
+
+@dataclass
+class A2CConfig:
+    env: EnvConfig = field(default_factory=EnvConfig)
+    train: TrainConfig = field(default_factory=TrainConfig)
+
+
+class A2CTrainer:
+    def __init__(self, config: A2CConfig) -> None:
+        self.config = config
+        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]
+        net = ActorCriticNet(env_dim=env_dim, action_num=action_num)
+        optimizer = optim.Adam(net.parameters(), lr=config.train.learning_rate)
+        self.agent = Agent(net=net, optimizer=optimizer)
+
+        self.num_episodes = config.train.num_episodes
+        self.gamma = config.train.gamma
+        self.lambda_ = config.train.lambda_
+        self.value_loss_coef = config.train.value_loss_coef
+        self.policy_loss_coef = config.train.policy_loss_coef
+        self.entropy_coef = config.train.entropy_coef
+        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="A2C",
+                name=f"[{config.env.env_name}]lr={config.train.learning_rate}",
+                # track hyperparameters and run metadata
+                config=asdict(config),
+            )
+
+    def train(self) -> None:
+        for i, episode in enumerate(range(self.num_episodes)):
+            observation, _ = self.env.reset()
+            terminated, truncated = False, False
+            while not (terminated or truncated):
+                action = self.agent.act(observation)
+                next_observation, reward, terminated, truncated, _ = self.env.step(action)
+                experience = Experience(
+                    observation=observation,
+                    action=action,
+                    reward=float(reward),
+                    terminated=terminated,
+                    truncated=truncated,
+                    next_observation=next_observation,
+                )
+                self.agent.add_experience(experience)
+                observation = next_observation
+                if self.config.env.render_mode is not None:
+                    self.env.render()
+            loss, action_entropy, advantages_mean = self.agent.net_update(
+                gamma=self.gamma,
+                lambda_=self.lambda_,
+                value_loss_coef=self.value_loss_coef,
+                policy_loss_coef=self.policy_loss_coef,
+                entropy_coef=self.entropy_coef,
+            )
+            total_reward = self.agent.get_buffer_total_reward()
+            solved = total_reward > self.solved_threshold
+            self.agent.onpolicy_reset()
+            print(
+                f"Episode {episode}, total_reward: {total_reward}, solved: {solved}, loss: {loss}, "
+                f"action_entropy: {action_entropy}, advantages_mean: {advantages_mean}"
+            )
+            if self.config.train.log_wandb:
+                wandb.log(
+                    {
+                        "episode": episode,
+                        "loss": loss,
+                        "total_reward": total_reward,
+                        "action_entropy": action_entropy,
+                        "advantages_mean": advantages_mean,
+                    }
+                )
+
+            if i % self.config.train.eval_interval == 0:
+                eval_reward = self.evaluate(self.config.train.eval_episodes)
+                print(f"Episode {episode}, Eval reward: {eval_reward}")
+                if self.config.train.log_wandb:
+                    wandb.log({"eval_reward": eval_reward, "episode": episode})
+
+    def evaluate(self, num_episodes: int) -> float:
+        total_reward = 0.0
+        for _ in range(num_episodes):
+            observation, _ = self.env.reset()
+            terminated, truncated = False, False
+            while not (terminated or truncated):
+                action = self.agent.act(observation, eval=True)
+                next_observation, reward, terminated, truncated, _ = self.env.step(action)
+                observation = next_observation
+                total_reward += float(reward)
+        return total_reward / num_episodes
+
+
+if __name__ == "__main__":
+    default_config = A2CConfig(
+        env=EnvConfig(env_name="CartPole-v1", render_mode=None, solved_threshold=475.0),
+        train=TrainConfig(
+            num_episodes=100000,
+            learning_rate=7e-4,
+            value_loss_coef=0.5,
+            policy_loss_coef=1,
+            entropy_coef=0.01,
+            log_wandb=True,
+        ),
+    )
+    trainer = A2CTrainer(default_config)
+    trainer.train()


Reviews:
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR introduces an Advantage Actor-Critic (A2C) reinforcement learning algorithm with configurable training and environment settings. Key changes include:
- A new neural network architecture (ActorCriticNet) with separate policy and value heads.
- Implementation of on-policy experience collection, Generalized Advantage Estimation (GAE), and logging support.
- Addition of training configuration and a command-line runnable training loop.




state: COMMENTED
author_association: NONE
submitted_at: 2025-05-03T06:15:25Z
user: login: coderabbitai[bot]
body: **Actionable comments posted: 5**

<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 2dd82a8fb32b731975563fe8ecbcbc955b36252a and cba7edcc55d7a56fcaa6630d7bd5eaf1cb452e25.

</details>

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

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

</details>

<details>
<summary>⏰ Context from checks skipped due to timeout of 90000ms (1)</summary>

* GitHub Check: ci (3.12, ubuntu-latest)

</details>

</details>

<!-- This is an auto-generated comment by CodeRabbit for review status -->
state: COMMENTED
author_association: NONE
submitted_at: 2025-05-03T06:17:14Z
user: login: coderabbitai[bot]
body: **Actionable comments posted: 0**

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

<details>
<summary>tests/test_a2c.py (4)</summary><blockquote>

`98-114`: **Make the test more environment-agnostic.**

The test has a comment that specifically mentions CartPole, which makes it less general for other environments.


```diff
-    assert action in range(action_num)  # For CartPole, should be 0 or 1
+    assert action in range(action_num)  # Action should be within valid range
```

---

`116-150`: **Consider using parameterized testing for config verification.**

The test verifies each configuration parameter individually, which is comprehensive but could be more maintainable with a parameterized approach.


You could refactor this to use dictionaries and loops for more maintainable assertions:

```python
def test_config():
    """Test the Config class."""
    config = A2CConfig()
    
    # Check default values
    expected_defaults = {
        "env.env_name": "CartPole-v1",
        "train.gamma": 0.999,
        "train.lambda_": 0.98,
        "train.value_loss_coef": 0.5,
        "train.policy_loss_coef": 0.5,
        "train.entropy_coef": 0.01,
    }
    
    for key, expected in expected_defaults.items():
        attr_path = key.split(".")
        actual = getattr(getattr(config, attr_path[0]), attr_path[1])
        assert actual == expected, f"Default {key} should be {expected}, got {actual}"
    
    # Test custom config
    custom_values = {
        "env.env_name": "MountainCar-v0",
        "env.solved_threshold": 90.0,
        "train.num_episodes": 1000,
        "train.learning_rate": 0.005,
        "train.gamma": 0.99,
        "train.lambda_": 0.95,
        "train.value_loss_coef": 0.7,
        "train.policy_loss_coef": 0.8,
        "train.entropy_coef": 0.02,
    }
    
    # Create custom config
    custom_config = A2CConfig(
        env=EnvConfig(env_name=custom_values["env.env_name"], 
                      solved_threshold=custom_values["env.solved_threshold"]),
        train=TrainConfig(
            num_episodes=custom_values["train.num_episodes"],
            learning_rate=custom_values["train.learning_rate"],
            gamma=custom_values["train.gamma"],
            lambda_=custom_values["train.lambda_"],
            value_loss_coef=custom_values["train.value_loss_coef"],
            policy_loss_coef=custom_values["train.policy_loss_coef"],
            entropy_coef=custom_values["train.entropy_coef"],
        ),
    )
    
    # Verify custom values
    for key, expected in custom_values.items():
        attr_path = key.split(".")
        actual = getattr(getattr(custom_config, attr_path[0]), attr_path[1])
        assert actual == expected, f"Custom {key} should be {expected}, got {actual}"
```

---

`173-213`: **Good test for the agent's learning step, consider adding more assertions.**

The test properly validates the agent's net_update method by creating experiences, calling the method, and checking the return types. 

Consider adding assertions to verify that the returned metrics are within reasonable ranges (non-NaN, finite values) and that the model parameters have changed after the update.


```diff
    # Test net_update
    loss, entropy, adv_mean = agent.net_update(
        gamma=0.99, lambda_=0.95, value_loss_coef=0.5, policy_loss_coef=0.5, entropy_coef=0.01
    )

    assert isinstance(loss, float)
    assert isinstance(entropy, float)
    assert isinstance(adv_mean, float)
+    # Check that values are finite (not NaN or infinite)
+    assert np.isfinite(loss)
+    assert np.isfinite(entropy)
+    assert np.isfinite(adv_mean)
+    
+    # Check that parameters were updated (optional)
+    # This would require capturing parameter values before the update
```

---

`215-230`: **Consider enhancing the minimal training test.**

While this test confirms basic functioning without errors, it doesn't verify that training actually improves performance. You might consider running a slightly longer training session and checking that rewards trend upward or that losses decrease.


```diff
def test_minimal_training():
    """Test minimal training run with a single episode."""
    # Create minimal config with just one episode
    config = A2CConfig(
        env=EnvConfig(env_name="CartPole-v1", render_mode=None),
-        train=TrainConfig(num_episodes=1, learning_rate=0.01, log_wandb=False),
+        train=TrainConfig(num_episodes=3, learning_rate=0.01, log_wandb=False),
    )

    # Initialize trainer
    trainer = A2CTrainer(config)

+    # Track metrics across episodes to verify learning is happening
+    rewards = []
+    losses = []
+    
+    # Override train method to collect metrics
+    original_train = trainer.train
+    def train_with_metrics():
+        for epi in range(trainer.num_episodes):
+            # Run a single episode
+            observation, _ = trainer.env.reset()
+            terminated, truncated = False, False
+            while not (terminated or truncated):
+                action, _ = trainer.agent.act(observation)
+                next_observation, reward, terminated, truncated, _ = trainer.env.step(action)
+                experience = Experience(
+                    observation=observation,
+                    action=action,
+                    reward=float(reward),
+                    terminated=terminated,
+                    truncated=truncated,
+                    next_observation=next_observation,
+                )
+                trainer.agent.add_experience(experience)
+                observation = next_observation
+            
+            # Get metrics
+            loss, _, _ = trainer.agent.net_update(
+                gamma=trainer.gamma,
+                lambda_=trainer.lambda_,
+                value_loss_coef=trainer.value_loss_coef,
+                policy_loss_coef=trainer.policy_loss_coef,
+                entropy_coef=trainer.entropy_coef,
+            )
+            total_reward = trainer.agent.get_buffer_total_reward()
+            rewards.append(total_reward)
+            losses.append(loss)
+            trainer.agent.onpolicy_reset()
+    
+    # Replace train method
+    trainer.train = train_with_metrics
+
    # Run training for one episode
    trainer.train()

    # If we got here without errors, test passed
    assert True
+    
+    # Check that metrics were collected
+    assert len(rewards) == 3
+    assert len(losses) == 3
+    
+    # Optionally check for improvement trends
+    # Note: This might be flaky with only 3 episodes
+    # assert rewards[-1] >= rewards[0], "Reward should not decrease"
```

Note: The above modification is optional and might make the test more flaky since RL training can be noisy. Consider if the benefit of deeper validation outweighs the potential for intermittent failures.

</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 cba7edcc55d7a56fcaa6630d7bd5eaf1cb452e25 and 59fd3c3895ca5f961abcca74b81c092c7b6583d2.

</details>

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

* `tests/test_a2c.py` (1 hunks)
* `toyrl/__init__.py` (1 hunks)

</details>

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

* toyrl/__init__.py

</details>

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

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

<details>
<summary>tests/test_a2c.py (1)</summary><blockquote>

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

* `A2CConfig` (151-153)
* `A2CTrainer` (156-224)
* `ActorCriticNet` (13-29)
* `net_update` (84-127)

</details>

</blockquote></details>

</details>

</details>

<details>
<summary>⏰ Context from checks skipped due to timeout of 90000ms (1)</summary>

* GitHub Check: ci (3.13, ubuntu-latest)

</details>

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

<details>
<summary>tests/test_a2c.py (5)</summary>

`1-16`: **Imports and file structure look good.**

The imports are appropriate and well-organized, containing all necessary packages (gymnasium, numpy, torch) and component imports from the A2C implementation.

---

`19-33`: **Thorough testing of neural network architecture.**

The test verifies the expected output shapes, types, and dtypes for both policy logits and value outputs of the ActorCriticNet forward pass. This ensures the network architecture produces outputs compatible with downstream components.

---

`36-81`: **Well-structured replay buffer testing.**

This test thoroughly validates the ReplayBuffer functionality through its lifecycle:
- Initial empty state checks
- Adding experiences
- Reward calculation
- Sampling behavior
- Reset functionality

Each operation's outcome is properly verified.

---

`153-170`: **Trainer creation test is thorough and well-structured.**

The test verifies that the trainer is properly initialized with the correct configuration parameters and that it has all the required components. The assertions are comprehensive.

---

`84-96`: **Agent creation test is concise and effective.**

The test properly verifies that the agent has all the expected attributes after initialization.

</details>

</blockquote></details>

</details>

<!-- This is an auto-generated comment by CodeRabbit for review status -->
state: COMMENTED
author_association: NONE
submitted_at: 2025-05-03T06:21:51Z

Review comments:
diff_hunk: @@ -0,0 +1,240 @@
+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 ActorCriticNet(nn.Module):
+    def __init__(self, env_dim: int, action_num: int) -> None:
+        super().__init__()
+        self.env_dim = env_dim
+        self.action_num = action_num
+        self.shared_layers = nn.Sequential(
+            nn.Linear(env_dim, 64),
+            nn.ReLU(),
+        )
+        self.policy_head = nn.Linear(64, action_num)
+        self.value_head = nn.Linear(64, 1)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        x = self.shared_layers(x)
+        policy_action_logits = self.policy_head(x)
+        v_value = self.value_head(x)
+        return policy_action_logits, v_value
+
+
+@dataclass
+class Experience:
+    observation: Any
+    action: Any
+    reward: float
+    next_observation: Any
+    terminated: bool
+    truncated: bool
+
+
+@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) -> list[Experience]:
+        return self.buffer
+
+    def total_reward(self) -> float:
+        return sum(experience.reward for experience in self.buffer)
+
+
+class Agent:
+    def __init__(self, net: nn.Module, optimizer: torch.optim.Optimizer) -> None:
+        self._net = net
+        self._optimizer = optimizer
+        self._replay_buffer = ReplayBuffer()
+
+    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) -> tuple[int, torch.Tensor]:
+        x = torch.from_numpy(observation.astype(np.float32))
+        action_logits, value = self._net(x)
+        next_action_dist = torch.distributions.Categorical(logits=action_logits)
+        action = next_action_dist.sample()
+        return action.item(), value
+
+    def net_update(
+        self, gamma: float, lambda_: float, value_loss_coef: float, policy_loss_coef: float, entropy_coef: float
+    ) -> tuple[float, float, float]:
+        experiences = self._replay_buffer.sample()
+
+        observations = torch.tensor([exp.observation for exp in experiences])
+        actions = torch.tensor([exp.action for exp in experiences])
+        rewards = torch.tensor([exp.reward for exp in experiences]).unsqueeze(1)
+        next_observations = torch.tensor([exp.next_observation for exp in experiences])
+        terminateds = torch.tensor([exp.terminated for exp in experiences], dtype=torch.float32).unsqueeze(1)
+
+        # calculate predicted V-values
+        policy_action_logits, v_values = self._net(observations)
+        # calculate advantages by GAE
+        with torch.no_grad():
+            _, v_values_next = self._net(next_observations)
+        deltas = rewards + gamma * v_values_next * (1 - terminateds) - v_values
+        advantages = deltas.clone()
+        for t in reversed(range(len(experiences) - 1)):
+            advantages[t] = deltas[t] + gamma * lambda_ * advantages[t + 1] * (1 - terminateds[t])
+        # note that advantages is not normalized here, we will normalize it in the policy loss
+        # this is because we want to keep the advantages as is in the value loss
+        v_targets = advantages + v_values
+        # calculate value loss
+        value_loss = nn.functional.mse_loss(v_values, v_targets)
+
+        # calculate policy loss
+        action_dist = torch.distributions.Categorical(logits=policy_action_logits)
+        action_entropy = action_dist.entropy().mean()
+        action_log_probs = action_dist.log_prob(actions)
+        # TODO: normalize advantages(has problem, disable for now)
+        # advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+        advantages = advantages / (advantages.std() + 1e-8)
path: toyrl/a2c.py
user: login: Copilot
body: Normalizing advantages by only dividing by the standard deviation omits centering; if full normalization is intended, subtract the mean before dividing by the standard deviation.
```suggestion
        # Normalize advantages by centering and scaling
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        # Full normalization ensures mean of zero and standard deviation of one
```
created_at: 2025-05-03T06:15:24Z
author_association: NONE
diff_hunk: @@ -0,0 +1,240 @@
+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 ActorCriticNet(nn.Module):
+    def __init__(self, env_dim: int, action_num: int) -> None:
+        super().__init__()
+        self.env_dim = env_dim
+        self.action_num = action_num
+        self.shared_layers = nn.Sequential(
+            nn.Linear(env_dim, 64),
+            nn.ReLU(),
+        )
+        self.policy_head = nn.Linear(64, action_num)
+        self.value_head = nn.Linear(64, 1)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        x = self.shared_layers(x)
+        policy_action_logits = self.policy_head(x)
+        v_value = self.value_head(x)
+        return policy_action_logits, v_value
+
+
+@dataclass
+class Experience:
+    observation: Any
+    action: Any
+    reward: float
+    next_observation: Any
+    terminated: bool
+    truncated: bool
+
+
+@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) -> list[Experience]:
+        return self.buffer
+
+    def total_reward(self) -> float:
+        return sum(experience.reward for experience in self.buffer)
+
+
+class Agent:
+    def __init__(self, net: nn.Module, optimizer: torch.optim.Optimizer) -> None:
+        self._net = net
+        self._optimizer = optimizer
+        self._replay_buffer = ReplayBuffer()
+
+    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) -> tuple[int, torch.Tensor]:
path: toyrl/a2c.py
user: login: Copilot
body: [nitpick] The type hint for 'observation' is specified as np.floating, which may be misleading; consider using np.ndarray to more accurately represent the input type.
```suggestion
    def act(self, observation: np.ndarray) -> tuple[int, torch.Tensor]:
```
created_at: 2025-05-03T06:15:24Z
author_association: NONE
diff_hunk: @@ -0,0 +1,240 @@
+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/a2c.py
user: login: coderabbitai[bot]
body: _🛠️ Refactor suggestion_

_⚠️ Potential issue_

**Make `wandb` an optional dependency**

`wandb` is imported unconditionally. On systems where it is not installed, the *entire script will crash* even when `log_wandb=False`.  Import it lazily or behind a `try/except` guard.

```diff
-import wandb
+try:
+    import wandb
+except ImportError:  # keep the library optional
+    wandb = None
```

Later, guard the `wandb.*` usages with `if config.train.log_wandb and wandb is not None:`.





<!-- 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
try:
    import wandb
except ImportError:  # keep the library optional
    wandb = None
`````

</details>

<!-- suggestion_end -->

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-05-03T06:17:13Z
author_association: NONE
diff_hunk: @@ -0,0 +1,240 @@
+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 ActorCriticNet(nn.Module):
+    def __init__(self, env_dim: int, action_num: int) -> None:
+        super().__init__()
+        self.env_dim = env_dim
+        self.action_num = action_num
+        self.shared_layers = nn.Sequential(
+            nn.Linear(env_dim, 64),
+            nn.ReLU(),
+        )
+        self.policy_head = nn.Linear(64, action_num)
+        self.value_head = nn.Linear(64, 1)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        x = self.shared_layers(x)
+        policy_action_logits = self.policy_head(x)
+        v_value = self.value_head(x)
+        return policy_action_logits, v_value
+
+
+@dataclass
+class Experience:
+    observation: Any
+    action: Any
+    reward: float
+    next_observation: Any
+    terminated: bool
+    truncated: bool
+
+
+@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) -> list[Experience]:
+        return self.buffer
+
+    def total_reward(self) -> float:
+        return sum(experience.reward for experience in self.buffer)
+
+
+class Agent:
+    def __init__(self, net: nn.Module, optimizer: torch.optim.Optimizer) -> None:
+        self._net = net
+        self._optimizer = optimizer
+        self._replay_buffer = ReplayBuffer()
+
+    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) -> tuple[int, torch.Tensor]:
+        x = torch.from_numpy(observation.astype(np.float32))
+        action_logits, value = self._net(x)
+        next_action_dist = torch.distributions.Categorical(logits=action_logits)
+        action = next_action_dist.sample()
+        return action.item(), value
+
+    def net_update(
+        self, gamma: float, lambda_: float, value_loss_coef: float, policy_loss_coef: float, entropy_coef: float
+    ) -> tuple[float, float, float]:
+        experiences = self._replay_buffer.sample()
+
+        observations = torch.tensor([exp.observation for exp in experiences])
+        actions = torch.tensor([exp.action for exp in experiences])
+        rewards = torch.tensor([exp.reward for exp in experiences]).unsqueeze(1)
+        next_observations = torch.tensor([exp.next_observation for exp in experiences])
+        terminateds = torch.tensor([exp.terminated for exp in experiences], dtype=torch.float32).unsqueeze(1)
+
+        # calculate predicted V-values
+        policy_action_logits, v_values = self._net(observations)
+        # calculate advantages by GAE
+        with torch.no_grad():
+            _, v_values_next = self._net(next_observations)
+        deltas = rewards + gamma * v_values_next * (1 - terminateds) - v_values
+        advantages = deltas.clone()
+        for t in reversed(range(len(experiences) - 1)):
+            advantages[t] = deltas[t] + gamma * lambda_ * advantages[t + 1] * (1 - terminateds[t])
+        # note that advantages is not normalized here, we will normalize it in the policy loss
path: toyrl/a2c.py
user: login: coderabbitai[bot]
body: _🛠️ Refactor suggestion_

**Include `truncated` episodes in the GAE bootstrapping**

Only `terminated` is considered when zeroing the bootstrapped value.  
Environments such as CartPole emit `truncated=True` when hitting the time-limit; those steps should also break the bootstrap chain.

```diff
-            advantages[t] = deltas[t] + gamma * lambda_ * advantages[t + 1] * (1 - terminateds[t])
+            done_t = done_flags[t]  # see previous comment
+            advantages[t] = deltas[t] + gamma * lambda_ * advantages[t + 1] * (1 - done_t)
```

Failing to do this *overestimates* advantages for time-limited terminations.




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

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-05-03T06:17:13Z
author_association: NONE
diff_hunk: @@ -0,0 +1,240 @@
+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 ActorCriticNet(nn.Module):
+    def __init__(self, env_dim: int, action_num: int) -> None:
+        super().__init__()
+        self.env_dim = env_dim
+        self.action_num = action_num
+        self.shared_layers = nn.Sequential(
+            nn.Linear(env_dim, 64),
+            nn.ReLU(),
+        )
+        self.policy_head = nn.Linear(64, action_num)
+        self.value_head = nn.Linear(64, 1)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        x = self.shared_layers(x)
+        policy_action_logits = self.policy_head(x)
+        v_value = self.value_head(x)
+        return policy_action_logits, v_value
+
+
+@dataclass
+class Experience:
+    observation: Any
+    action: Any
+    reward: float
+    next_observation: Any
+    terminated: bool
+    truncated: bool
+
+
+@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) -> list[Experience]:
+        return self.buffer
+
+    def total_reward(self) -> float:
+        return sum(experience.reward for experience in self.buffer)
+
+
+class Agent:
+    def __init__(self, net: nn.Module, optimizer: torch.optim.Optimizer) -> None:
+        self._net = net
+        self._optimizer = optimizer
+        self._replay_buffer = ReplayBuffer()
+
+    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) -> tuple[int, torch.Tensor]:
+        x = torch.from_numpy(observation.astype(np.float32))
+        action_logits, value = self._net(x)
+        next_action_dist = torch.distributions.Categorical(logits=action_logits)
+        action = next_action_dist.sample()
+        return action.item(), value
path: toyrl/a2c.py
user: login: coderabbitai[bot]
body: _⚠️ Potential issue_

**Disable gradient tracking and keep tensors on the correct device in `act`**

1. Action selection is inference-only; tracking gradients here wastes memory.
2. The created tensor is **not moved to the same device** as the network (e.g., GPU).

```diff
-    def act(self, observation: np.floating) -> tuple[int, torch.Tensor]:
-        x = torch.from_numpy(observation.astype(np.float32))
-        action_logits, value = self._net(x)
+    def act(self, observation: np.ndarray) -> tuple[int, torch.Tensor]:
+        with torch.no_grad():
+            x = torch.from_numpy(observation.astype(np.float32)).to(next(self._net.parameters()).device)
+            action_logits, value = self._net(x)
```

This avoids building needless graphs and prevents device mismatch errors.




<!-- 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.ndarray) -> tuple[int, torch.Tensor]:
        with torch.no_grad():
            x = torch.from_numpy(observation.astype(np.float32)).to(
                next(self._net.parameters()).device
            )
            action_logits, value = self._net(x)
        next_action_dist = torch.distributions.Categorical(logits=action_logits)
        action = next_action_dist.sample()
        return action.item(), value
`````

</details>

<!-- suggestion_end -->

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-05-03T06:17:13Z
author_association: NONE
diff_hunk: @@ -0,0 +1,240 @@
+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 ActorCriticNet(nn.Module):
+    def __init__(self, env_dim: int, action_num: int) -> None:
+        super().__init__()
+        self.env_dim = env_dim
+        self.action_num = action_num
+        self.shared_layers = nn.Sequential(
+            nn.Linear(env_dim, 64),
+            nn.ReLU(),
+        )
+        self.policy_head = nn.Linear(64, action_num)
+        self.value_head = nn.Linear(64, 1)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        x = self.shared_layers(x)
+        policy_action_logits = self.policy_head(x)
+        v_value = self.value_head(x)
+        return policy_action_logits, v_value
+
+
+@dataclass
+class Experience:
+    observation: Any
+    action: Any
+    reward: float
+    next_observation: Any
+    terminated: bool
+    truncated: bool
+
+
+@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) -> list[Experience]:
+        return self.buffer
+
+    def total_reward(self) -> float:
+        return sum(experience.reward for experience in self.buffer)
+
+
+class Agent:
+    def __init__(self, net: nn.Module, optimizer: torch.optim.Optimizer) -> None:
+        self._net = net
+        self._optimizer = optimizer
+        self._replay_buffer = ReplayBuffer()
+
+    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) -> tuple[int, torch.Tensor]:
+        x = torch.from_numpy(observation.astype(np.float32))
+        action_logits, value = self._net(x)
+        next_action_dist = torch.distributions.Categorical(logits=action_logits)
+        action = next_action_dist.sample()
+        return action.item(), value
+
+    def net_update(
+        self, gamma: float, lambda_: float, value_loss_coef: float, policy_loss_coef: float, entropy_coef: float
+    ) -> tuple[float, float, float]:
+        experiences = self._replay_buffer.sample()
+
+        observations = torch.tensor([exp.observation for exp in experiences])
+        actions = torch.tensor([exp.action for exp in experiences])
+        rewards = torch.tensor([exp.reward for exp in experiences]).unsqueeze(1)
+        next_observations = torch.tensor([exp.next_observation for exp in experiences])
+        terminateds = torch.tensor([exp.terminated for exp in experiences], dtype=torch.float32).unsqueeze(1)
+
+        # calculate predicted V-values
+        policy_action_logits, v_values = self._net(observations)
+        # calculate advantages by GAE
+        with torch.no_grad():
+            _, v_values_next = self._net(next_observations)
+        deltas = rewards + gamma * v_values_next * (1 - terminateds) - v_values
+        advantages = deltas.clone()
+        for t in reversed(range(len(experiences) - 1)):
+            advantages[t] = deltas[t] + gamma * lambda_ * advantages[t + 1] * (1 - terminateds[t])
+        # note that advantages is not normalized here, we will normalize it in the policy loss
+        # this is because we want to keep the advantages as is in the value loss
+        v_targets = advantages + v_values
+        # calculate value loss
+        value_loss = nn.functional.mse_loss(v_values, v_targets)
+
+        # calculate policy loss
+        action_dist = torch.distributions.Categorical(logits=policy_action_logits)
+        action_entropy = action_dist.entropy().mean()
+        action_log_probs = action_dist.log_prob(actions)
+        # TODO: normalize advantages(has problem, disable for now)
+        # advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+        advantages = advantages / (advantages.std() + 1e-8)
+        policy_loss = -action_log_probs * advantages
+        policy_loss = torch.mean(policy_loss)
path: toyrl/a2c.py
user: login: coderabbitai[bot]
body: _⚠️ Potential issue_

**Detach advantage & target tensors to prevent incorrect gradient flow**

`advantages` and `v_targets` contain `v_values`, which has gradients.  
Using them directly in the policy loss causes **second-order gradient leakage** into the value network, destabilising learning.

```diff
-        v_targets = advantages + v_values
+        advantages = advantages.detach()
+        v_targets = (advantages + v_values).detach()
```

If you need `advantages` *un-detached* for the value loss, create two copies—one detached for the policy loss, one kept for the value computation.




<!-- 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
        deltas = rewards + gamma * v_values_next * (1 - terminateds) - v_values
        advantages = deltas.clone()
        for t in reversed(range(len(experiences) - 1)):
            advantages[t] = deltas[t] + gamma * lambda_ * advantages[t + 1] * (1 - terminateds[t])
        # note that advantages is not normalized here, we will normalize it in the policy loss
        # this is because we want to keep the advantages as is in the value loss
        advantages = advantages.detach()
        v_targets = (advantages + v_values).detach()
        # calculate value loss
        value_loss = nn.functional.mse_loss(v_values, v_targets)

        # calculate policy loss
        action_dist = torch.distributions.Categorical(logits=policy_action_logits)
        action_entropy = action_dist.entropy().mean()
        action_log_probs = action_dist.log_prob(actions)
        # TODO: normalize advantages(has problem, disable for now)
        # advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        advantages = advantages / (advantages.std() + 1e-8)
        policy_loss = -action_log_probs * advantages
        policy_loss = torch.mean(policy_loss)
`````

</details>

<!-- suggestion_end -->

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-05-03T06:17:13Z
author_association: NONE
diff_hunk: @@ -0,0 +1,240 @@
+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 ActorCriticNet(nn.Module):
+    def __init__(self, env_dim: int, action_num: int) -> None:
+        super().__init__()
+        self.env_dim = env_dim
+        self.action_num = action_num
+        self.shared_layers = nn.Sequential(
+            nn.Linear(env_dim, 64),
+            nn.ReLU(),
+        )
+        self.policy_head = nn.Linear(64, action_num)
+        self.value_head = nn.Linear(64, 1)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        x = self.shared_layers(x)
+        policy_action_logits = self.policy_head(x)
+        v_value = self.value_head(x)
+        return policy_action_logits, v_value
+
+
+@dataclass
+class Experience:
+    observation: Any
+    action: Any
+    reward: float
+    next_observation: Any
+    terminated: bool
+    truncated: bool
+
+
+@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) -> list[Experience]:
+        return self.buffer
+
+    def total_reward(self) -> float:
+        return sum(experience.reward for experience in self.buffer)
+
+
+class Agent:
+    def __init__(self, net: nn.Module, optimizer: torch.optim.Optimizer) -> None:
+        self._net = net
+        self._optimizer = optimizer
+        self._replay_buffer = ReplayBuffer()
+
+    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) -> tuple[int, torch.Tensor]:
+        x = torch.from_numpy(observation.astype(np.float32))
+        action_logits, value = self._net(x)
+        next_action_dist = torch.distributions.Categorical(logits=action_logits)
+        action = next_action_dist.sample()
+        return action.item(), value
+
+    def net_update(
+        self, gamma: float, lambda_: float, value_loss_coef: float, policy_loss_coef: float, entropy_coef: float
+    ) -> tuple[float, float, float]:
+        experiences = self._replay_buffer.sample()
+
+        observations = torch.tensor([exp.observation for exp in experiences])
+        actions = torch.tensor([exp.action for exp in experiences])
+        rewards = torch.tensor([exp.reward for exp in experiences]).unsqueeze(1)
+        next_observations = torch.tensor([exp.next_observation for exp in experiences])
+        terminateds = torch.tensor([exp.terminated for exp in experiences], dtype=torch.float32).unsqueeze(1)
+
path: toyrl/a2c.py
user: login: coderabbitai[bot]
body: _💡 Verification agent_

<details>
<summary>🧩 Analysis chain</summary>

**Explicitly set dtype & device when batching experience tensors**

`torch.tensor([...])` defaults to `float64` on CPU, whereas the network is built with `float32`.  
This silent mismatch slows computation and may break GPU execution.  
Additionally, the `truncated` flag is ignored.

```diff
-        observations = torch.tensor([exp.observation for exp in experiences])
-        actions = torch.tensor([exp.action for exp in experiences])
-        rewards = torch.tensor([exp.reward for exp in experiences]).unsqueeze(1)
-        next_observations = torch.tensor([exp.next_observation for exp in experiences])
-        terminateds = torch.tensor([exp.terminated for exp in experiences], dtype=torch.float32).unsqueeze(1)
+        device = next(self._net.parameters()).device
+        observations = torch.as_tensor(
+            [exp.observation for exp in experiences], dtype=torch.float32, device=device
+        )
+        actions = torch.as_tensor(
+            [exp.action for exp in experiences], dtype=torch.int64, device=device
+        )
+        rewards = torch.as_tensor(
+            [exp.reward for exp in experiences], dtype=torch.float32, device=device
+        ).unsqueeze(1)
+        next_observations = torch.as_tensor(
+            [exp.next_observation for exp in experiences], dtype=torch.float32, device=device
+        )
+        done_flags = torch.as_tensor(
+            [exp.terminated or exp.truncated for exp in experiences], dtype=torch.float32, device=device
+        ).unsqueeze(1)
```

(Use `done_flags` in place of `terminateds` later.)

---

🏁 Script executed:

```shell
#!/bin/bash
set -e
# Locate the file and show relevant lines for context
file=$(find . -type f | grep "toyrl/a2c.py" || true)
echo "File path: $file"
echo "--- Lines 1-200 of toyrl/a2c.py ---"
sed -n '1,200p' "$file"

# Check for Experience class definition
echo "--- Searching for Experience class ---"
rg -n "class .*Experience" -n .

# Check for any use of 'truncated' in the repo
echo "--- Searching for 'trunc' occurrences ---"
rg -n "trunc" -n .
```

Length of output: 9666

---



</details>

**Explicitly specify tensor dtype & device and include `truncated` in done‐mask**

The current batching uses plain `torch.tensor(…)` (inference from data) and only `terminated`, which can lead to silent dtype/device mismatches and treats truncated episodes as continuing. Please update in `toyrl/a2c.py` inside `Agent.net_update` (around lines 89–94):

- Move tensors to the network’s device and set a consistent dtype (e.g. `float32` for floats, `int64` for actions)
- Combine `terminated` OR `truncated` into a single done‐flag

Suggested diff:
```diff
-        observations = torch.tensor([exp.observation for exp in experiences])
-        actions = torch.tensor([exp.action for exp in experiences])
-        rewards = torch.tensor([exp.reward for exp in experiences]).unsqueeze(1)
-        next_observations = torch.tensor([exp.next_observation for exp in experiences])
-        terminateds = torch.tensor([exp.terminated for exp in experiences], dtype=torch.float32).unsqueeze(1)
+        device = next(self._net.parameters()).device
+        observations = torch.as_tensor(
+            [exp.observation for exp in experiences],
+            dtype=torch.float32, device=device
+        )
+        actions = torch.as_tensor(
+            [exp.action for exp in experiences],
+            dtype=torch.int64, device=device
+        )
+        rewards = torch.as_tensor(
+            [exp.reward for exp in experiences],
+            dtype=torch.float32, device=device
+        ).unsqueeze(1)
+        next_observations = torch.as_tensor(
+            [exp.next_observation for exp in experiences],
+            dtype=torch.float32, device=device
+        )
+        done_flags = torch.as_tensor(
+            [exp.terminated or exp.truncated for exp in experiences],
+            dtype=torch.float32, device=device
+        ).unsqueeze(1)
```
Use `done_flags` in place of `terminateds` throughout the advantage and mask calculations.  


<!-- 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
        device = next(self._net.parameters()).device
        observations = torch.as_tensor(
            [exp.observation for exp in experiences],
            dtype=torch.float32, device=device
        )
        actions = torch.as_tensor(
            [exp.action for exp in experiences],
            dtype=torch.int64, device=device
        )
        rewards = torch.as_tensor(
            [exp.reward for exp in experiences],
            dtype=torch.float32, device=device
        ).unsqueeze(1)
        next_observations = torch.as_tensor(
            [exp.next_observation for exp in experiences],
            dtype=torch.float32, device=device
        )
        done_flags = torch.as_tensor(
            [exp.terminated or exp.truncated for exp in experiences],
            dtype=torch.float32, device=device
        ).unsqueeze(1)
`````

</details>

<!-- suggestion_end -->

<!-- This is an auto-generated comment by CodeRabbit -->
created_at: 2025-05-03T06:17:13Z
author_association: NONE

PR comments:
user: login: coderabbitai[bot]
created_at: 2025-05-03T06:12:51Z
body: <!-- This is an auto-generated comment: summarize by coderabbit.ai -->
<!-- walkthrough_start -->

"""
## Walkthrough

A new file has been added that implements a complete Advantage Actor-Critic (A2C) reinforcement learning pipeline using PyTorch and Gymnasium. The code introduces neural network, agent, buffer, configuration, and trainer classes, as well as experience dataclasses. It supports on-policy experience collection, Generalized Advantage Estimation (GAE), entropy regularization, and optional logging to Weights & Biases. The script includes a main entry point to configure and run training.

## Changes

| File(s)           | Change Summary                                                                                                                                                                                                                                  |
|-------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| toyrl/a2c.py      | Added a full A2C RL implementation: includes `ActorCriticNet`, `Experience`, `ReplayBuffer`, `Agent`, `EnvConfig`, `TrainConfig`, `A2CConfig`, `A2CTrainer`, and a main script entry point for configuration and training execution.            |
| tests/test_a2c.py | Added comprehensive tests for A2C components: neural network forward pass, replay buffer operations, agent initialization and action selection, configuration defaults, trainer instantiation, minimal training run, and evaluation method.  |
| toyrl/__init__.py | Extended imports and `__all__` exports to include `A2CConfig` and `A2CTrainer` from the `a2c` submodule.                                                                                                                                         |

## Sequence Diagram(s)

```mermaid
sequenceDiagram
    participant Main
    participant Trainer
    participant Agent
    participant Env
    participant Buffer
    participant Net

    Main->>Trainer: Instantiate with config
    Trainer->>Env: Initialize environment
    Trainer->>Agent: Initialize with Net, optimizer, buffer
    loop For each episode
        Trainer->>Env: Reset environment
        loop Until episode ends
            Agent->>Net: Forward state
            Net-->>Agent: Action logits, value
            Agent->>Env: Select and perform action
            Env-->>Agent: Next state, reward, done
            Agent->>Buffer: Store experience
        end
        Agent->>Net: Update parameters using GAE
        Net-->>Agent: Losses, backpropagation
        Trainer->>Main: Log progress
    end
```

## Assessment against linked issues

| Objective                                                                                      | Addressed | Explanation |
|------------------------------------------------------------------------------------------------|-----------|-------------|
| Implement A2C algorithm with policy and value networks, experience collection, and training loop (#23) | ✅        |             |
| Use Generalized Advantage Estimation (GAE) for advantage calculation (#23)                      | ✅        |             |
| Provide configurable environment and training hyperparameters (#23)                            | ✅        |             |
| Include logging support for training progress (#23)                                           | ✅        |             |

## Poem

> In the land of code, a new path unfurled,  
> A2C hops in, to the RL world.  
> Actors and critics, a clever pair,  
> Learn from rewards with rabbit care.  
> Buffers and trainers, configs in tow,  
> Watch your agents learn and grow!  
> 🐇✨
"""

<!-- walkthrough_end -->


---

<details>
<summary>📜 Recent 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 89ad896db2f1468a1ec592a997f8081d496157cf and b5c678ff6851da665ea2854619b03f44c16e72f8.

</details>

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

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

</details>

<details>
<summary>🚧 Files skipped from review as they are similar to previous changes (1)</summary>

* toyrl/a2c.py

</details>

<details>
<summary>⏰ Context from checks skipped due to timeout of 90000ms (2)</summary>

* GitHub Check: Docs (python-3.13, ubuntu-latest)
* GitHub Check: ci (3.12, ubuntu-latest)

</details>

</details>
<!-- internal state start -->


<!-- DwQgtGAEAqAWCWBnSTIEMB26CuAXA9mAOYCmGJATmriQCaQDG+Ats2bgFyQAOFk+AIwBWJBrngA3EsgEBPRvlqU0AgfFwA6NPEgQAfACgjoCEYDEZyAAUASpETZWaCrKNxU3bABsvkCiQBHbGlcABpIcVwvOkgAIgAzEmouNFp6AEEAJgBhUNjIAHc0ZAcBZnUaejkI2BJIbERKe1qMAA94TCIAL1hsTvRkTx8/QODEXEgMRwEmzMyULFxake58RHV8Fw0YZaVEBgp4bnF8LH9E/wwGaUYvNfgMIhREBzqp5hm+TIBmcIfaeAMagPJ5Lag1OoMWCdG4PXAURTYa70B41VC2dBpfwvGrgpCvbYAOXwkHi2AoSyaSlw2i8yBU+DwEJQzG40TYGBpJywm1JSVw5JuzjqvHwEngShRi120gOR25GnMFkgAFVuLRqDFBCIxJIbg4nC43OpolxEslMfQ0JkGJBIAZCdNKFw5naDCrGhQuIgWu1Oj0+o8jPpjOAoGR6Ph4jgCMQyMpKgpWOwuLx+MJROIpDJ5EwlFRVOotDoQyYoO5kKhMDHCKRyFRE0xk5yuFQCvZHMxnPJqnnlIXNNpdGBDKHTAYCLIKF4APTWhgabiyDgGWJrgzK9IASTj9c19ANXZc/GjUJhiDcCGQ5Hb8Xg0RZbJIHNw9Kw6VoEkwNNIkHSYk2MBskOcRbQACiybIAEoRgeeJNmuF9IGiZwMBBCIqAedDuCOEgvAeOoGnQqxZGgBDYHQDB6AAcVkZgMGKeBHG2LcJiUO9yDfSAAAN/wIChgPUQFCRIXBuMmEhyTQXxyFwApNgAa0KdQKJ9YV6C8NBZEoN8DxIbhnE1Hh8Hwhh5Ewegvy8YJIFqVJkAIfg8E8CY0F1U5kPwIh1F0+wuTqKybJCeAuxoRBtgAUVabhKHgMhrgibAn3pfw/M2GJUTQHibH0zTZAAIWweILnEhhNJecICgQKEO24VYKXpNIQXCbFRPER5wgs+w0FZfDHkgEhoti+LpG2OA6l4usxNuYpkC7BjSAc5ZZPkigFPCfBjhC+AukoTqqJWPLIAEIqLnCUVxQBfq2CWRQHJJNyJlRSkBowcUEQwJCuqSjUaGZFbFJ4Qybp0+p1n66j4yofDdoyT9vzQX8IvGELgQ8sDqPSCKoLG5YfqMpsXJuQK6juHFxSythq0QIJ1IGigET4GY5JIMgeH8AExBiLqaQoUgJhJxBzpMwF5DJ5AiP6l67ieUUBBUe8hJuAoSHgIhYETapUi/TlEekfarSwdgESXY7TgaCISXixkqF/Qa2U2NGMFx0m1kaFLIRYNRyENqo3IU0UDKIfdLfqdUjJegG1sVAxslOO8iCk7lIF+tzyvdyAwO4iK3rjjAE+48JuOgTCMDzguDZ4yDy7V7iYPitBuAcTS/rId7Ti+g74W0ND+tgWQYooAyqBBihwp2CbIJLnvKFK9OFiE6SdpuF62/gD6X3CKP1v4Lbyl2ihK71zlK4obAMCWupu6w/q7k20leSyxAYoYeA7xid5PhPAacMQRRRsgBFNyFF9JIH/goHwmZkD22GlcG459xC+BoBQcoDFk68nhOfIE3JwiUiwIPeCKD6QSTkoDfGNAxql2wgiIg2JKxXGsldUE+AaS+H8EUCgtBwji3CMbTasgj7w11r+KmWAPCHE5NzA6m1uTSS8GLLypB6COQAOqq3Vq+SAAAySA+UOju22FYWKihAQDUCk7Pw58PY8EoIQtgVoJhMHzmrQUUpkFWSFspJY6ApC2zqOw5wtBkAyyUZRegvA4R0BdvYOUxwFhlWwHsdAkAuyomNvIVYcJcQTEaJo9iaBvCOPji4hs8BTi+XGM4TR19e5EBjgAeV8XI3Byw+zszFBKIUSYnw0HCMwREmlD5hIUM4xOBYHwhSfC+CxUZmSQXQF4IgmwVLMC8RRU4YBVimXkDAw4I0IHRHchgcImMIpgHlo0K0QifyQmkgwbwTs/iSNoSHEg5icHDJCehBwdVNiaCMEqP8XhkFO3usyJQZVDLcmQLM+2fytR8E8AIUyr1IhxQvAYKAH5JQzRxLxACAkQLCVEuJVE3FJzTjnDaRcshxIFBUiM8YZ8CXDO4oQjhtBxI3VgIoRUWK0gxFTpCvFUVB5xTgaSrA5L8BTlnPOGl4l/C8GkOwdCT8QQPm7hfISpw+V/gFfQIV89uI5TZFpQqxVZ4LB4hSuV1Klx0oZdyu6PEAD6rrogYHdYXHiqRaCut2eK64PruKtTEkXRAPUnw+q6tK1hrqAmcO4nq7FMRhXIEmuwSVNqZWUvlQ6tZKTRI8qCW611WFcDeqLqcLZosE0qvDb6tIAahp7IlUXfmrqTqWooK6gg8bE2cqLo9GNB1uKyVdeQkgybMX6pxUa2aPEc4SBrkQbN0rZVUoXAW+l3j35eFLdxNurqGJsBDZcfMrr+lKFHfQbif8vBSH9UsbEPKD0zv5fO6gadF3F1Lqu9dtqt0KsLfu0thCMI93Qv3Qew8eqiVBg8BJTCeLBycCGzSHwNSup9eLBQJBiqAnFa+LeTo+Cwt/v/TxKEKC1L8JqXh7yyn4OBghselcvn9XiJpOps7U2Gu/emquOQAPWo3Xm+1tKkze3QtnXOxS12sung8ADKaDW4ozVPUuVqyVAfzVJ3dFEnFMqRPxIuNTb08TedJPoNAuXFrumpnFqSsD7EOHE9Jxksnn3zDxN+kB3WnpIO6yAABeULcR3UufdbEDggHc12u3VJxyDxKmcg6H9XiORlP1nEl1M+iwqGPEVGuWIwZxxhVfDOSrrr9MrlKxuSw25dwJhiIebs38zyPGkJeVAN4IghBSQMkUCJLo3EJv4Fo6wpD1DQq5A6kSXnJ0q8gCDL15kQWuXrP8BKgJEoYDBfwcEELPnYMhJItH0KTPZOwJ2LEim+OQEwVKhNTjsHoch9C62CWCVAiJCY5ApIyVEqtHepq8oWouA/Pggb9njEdqQcI6QpqvSBE3B57VZYizMp8i7tTwhOITknZjGn9bDPW9l7TfBsSrG1ciuoa2iugmWMfJ6WBqxrw3uwGO40BvjGQL4t+8gXrssCUDHEszvv8V+8SsIr0HASKeIyXALlmiNyFF3AeNwIO1pxzLHywySbszoICaF0TwfmtOk0FbkArISiMn69CsO4GeMJngL7LDpIjA5Z4yNvVmrDLDZj0kWDZH4VwLIbYyOzs28J+vVZgcmgVo6DDR56Ajn2DwpmEnopaBInQnbw1SADgIfT9CyukdRKTvDn9Z1yjoQTAIZsZgxCaN0fGPpEYApLu3zdrwzkJsBHDO1gjERSQsA3UOAwZAVV7wjZVRQcUUtlhKqOt2i4MdV3E48rH2oDAlL5MKcM+58OE+seQbbmz2veSc47mdnmTPM+4Ex+PHL1uQhPeKSghQDNMwLDS+IDMtGLfp9PfgdKzhXg3pBjfE8ECGcOfCklhF2D4PIBBm9j/GAkoGssrvTIzOPNijqgxCgZ1BJO2JVrbrFO/JfNAfWAAOTQJMYeR14E574KRfZQHWbWQWL+Dd4XzJLcb4DgiDprJpJDRcxVCMhUThRGCNLKA+AtJXwf4DStCUCvyNDMguYjLwgmSkh3DtgGRLCeKpy2QWR9RECVzK6q7Pb+BiCcQwrRgU7ZCPg3a6zJyvbkCcieJkAK5O5iGawkDQjiiPwLbPKlK6oAqNZAogrQqhwvQQqDKgrfxwoUgIo8DYDIqmKqriA9azoiRkGDZkhXAxF+oZRSorbVYhC1aSbcQrh2hQDkqVGPSbCupyigQnqiRgRQTcQGB1E2qVGr5aRdpW4UCdHdG9ENHjC1ZTQtH+BOyjE9G6B9GTGs61ZiDzHjE1Zx5EDrGLETGVo1KUAzH8jMY7H1E1YrETpTqnFLGVqoKoxeB9pM7XF7GPEzy9qcG2YnFdGrjriYoVYJYzjuoVruo0r1bribg7h1itYHidgdazJdaLSXh1CTJ/J+Sagvj0IaHDaFCUB1CDQ0BUQxApYMKJKTwiYKZ5ZjpaZvHiTxAIirIvTcTzjiSlDXreAkDbBxy/7PynBMLyIKFlpyLerIRIATBFASw15ElWzRSomUjqGkKkGopKyeJdhsHL7PhpEZFlTmQMDXAvDwD06FAtDMgGT77baoAokpG0BEgkgsK1B8D662i8hOLaG+ACHtgInKy4kpKpAckRGbjAoJjMZgpxGiAJExGwoylWn8CIrpEopZHorBhzoxDJETCzJZbZCiaxrUkEQUDZqMnupCk4YinjBiZAZAlzYgkOo/FlZ/GjhlivSRjRgFKxhQkNhposAvithoDtjtbHi9j/wFhqCDglgjgGANlNjlCVoSiID1rigkAqz+qVIUjDj1lhiMDywADsdAupAArLubQJuWgLuQAGzxBAhoAnknnfAAAMh5AgtAu5SQ8QAAjAwAIAACy7mZAkCZC7mrnjnrm7kACc8QtA3wDA3wAAHMBbuUCLufEMBSeS+SoLqWgJuR+QIFBW+TecBTaJuQICebuVBd8LQPMKWOuSeXhShR+S+TefEOhfLHeVBWgMBTebuW+ZkJuZkPECxSQC+VBYRQIHeWxURQBQ2TeSeR+d8AIN8HRd8LeXhVBVBQwC+cBQwLhS+ZucBcJZkIRaRd8KID+SQB+SeTeeJeuTBakDBSebQAIDxS+aZSxS+aICBZkKxcBZubxTedhbQB+UhS+buZuQwNGBRRAMdHBSeZuVBcVCeVBRxRqFeU+daPFaZWpcJd8PEB+R+apSeSQNxbxQBYBeFZOeoOWkEnOXFIuQGgdPoEAA=== -->

<!-- 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=22):

- 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.`

### Support

Need help? Create a ticket on our [support page](https://www.coderabbit.ai/contact-us/support) for assistance with any issues or questions.

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-05-03T06:20:39Z
body: ## [Codecov](https://app.codecov.io/gh/ai-glimpse/toyrl/pull/22?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ai-glimpse) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 98.23%. Comparing base [(`2dd82a8`)](https://app.codecov.io/gh/ai-glimpse/toyrl/commit/2dd82a8fb32b731975563fe8ecbcbc955b36252a?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ai-glimpse) to head [(`b5c678f`)](https://app.codecov.io/gh/ai-glimpse/toyrl/commit/b5c678ff6851da665ea2854619b03f44c16e72f8?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.

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


```diff
@@            Coverage Diff             @@
##           master      #22      +/-   ##
==========================================
+ Coverage   97.40%   98.23%   +0.83%     
==========================================
  Files           4        5       +1     
  Lines         193      284      +91     
==========================================
+ Hits          188      279      +91     
  Misses          5        5              
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/ai-glimpse/toyrl/pull/22?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