drlab 0.1.1__tar.gz → 0.2.2__tar.gz

drlab-0.2.2/PKG-INFO

@@ -0,0 +1,208 @@
+Metadata-Version: 2.4
+Name: drlab
+Version: 0.2.2
+Summary: Deep Reinforcement Learning kit for research.
+Author: Tomas Osarte
+Keywords: reinforcement-learning,deep-learning,pytorch,gymnasium
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch<3,>=2.0
+Requires-Dist: numpy==2.4.2
+Requires-Dist: gymnasium==1.2.3
+Requires-Dist: tqdm>=4.66
+Requires-Dist: tensorboard>=2.16
+Provides-Extra: experiments
+Requires-Dist: ale-py>=0.11; extra == "experiments"
+Requires-Dist: gymnasium[box2d]==1.2.3; extra == "experiments"
+Requires-Dist: matplotlib>=3.8; extra == "experiments"
+Requires-Dist: minigrid>=3.0; extra == "experiments"
+Requires-Dist: opencv-python>=4.10; extra == "experiments"
+Requires-Dist: scipy>=1.11; extra == "experiments"
+Provides-Extra: dev
+Requires-Dist: build>=1.2; extra == "dev"
+Requires-Dist: ipykernel>=6.29; extra == "dev"
+Requires-Dist: notebook>=7.0; extra == "dev"
+Requires-Dist: pip-tools>=7.4; extra == "dev"
+Requires-Dist: pytest>=8.0; extra == "dev"
+Requires-Dist: ruff>=0.6; extra == "dev"
+Dynamic: license-file
+
+# drlab
+
+`drlab` is a small deep reinforcement learning library for training neural
+network agents in Gymnasium environments. Its purpose is to provide simple,
+research-friendly building blocks for reinforcement learning experiments,
+especially when the environment or training setup is still being explored.
+
+The library is built around PyTorch models, Gymnasium environments, reusable
+controllers, replay buffers, runners, learners, and lightweight experiment
+loops. It is intended for research code and prototypes, not as a large
+production RL framework.
+
+## Installation
+
+The intended development setup uses [`uv`](https://docs.astral.sh/uv/). From the
+repository root:
+
+```bash
+uv sync --extra experiments --extra dev
+```
+
+Run commands through the managed environment:
+
+```bash
+uv run python -m pytest
+uv run python examples/dqn_cartpole.py
+```
+
+You can also activate the local virtual environment created by `uv`:
+
+```bash
+source .venv/bin/activate
+python -m pytest
+```
+
+If you prefer a normal `venv` and `pip` workflow:
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+python -m pip install --upgrade pip
+python -m pip install -e .
+```
+
+For experiment and development dependencies:
+
+```bash
+python -m pip install -e ".[experiments,dev]"
+```
+
+`drlab` depends on PyTorch, NumPy, Gymnasium, tqdm, and TensorBoard. If you need
+a specific CUDA build of PyTorch, install the matching PyTorch wheel before
+installing the package dependencies.
+
+## Library Overview
+
+Most public classes can be imported directly from `drlab`:
+
+```python
+from drlab import DQNConfig, DQNLearner, OffPolicyExperiment
+```
+
+### Off-Policy Learning
+
+Off-policy learners train from transition batches that may come from a replay
+buffer instead of only the most recent rollout. `OffPolicyExperiment` collects
+environment transitions with a `Runner`, stores them in a `ReplayBuffer`, samples
+minibatches, and calls the learner update.
+
+Implemented off-policy learners:
+
+- `DQNLearner`: value-based learning for discrete action spaces. It supports
+  replay-buffer training, target networks, Double DQN, hard or soft target
+  updates, gradient clipping, and custom regularizers.
+- `SACLearner`: Soft Actor-Critic style learning for continuous action spaces.
+  It uses an actor, two critics, target critics, entropy tuning, and replay
+  buffer training.
+
+### On-Policy Learning
+
+On-policy learners train from the data collected by the current policy.
+`OnPolicyExperiment` collects rollouts, computes returns when needed, and trains
+the learner directly on the fresh batch.
+
+Implemented on-policy learners:
+
+- `ReinforceLearner`: policy-gradient learning from discounted returns.
+- `ActorCriticLearner`: policy and value learning from shared model outputs.
+- `PPOLearner`: clipped policy optimization over rollout batches.
+
+### Supporting Components
+
+Controllers turn model outputs into environment actions:
+
+- `GreedyController`, `EpsilonGreedyController`, and `StochasticController`
+  support discrete action spaces.
+- `GaussianController` supports continuous action spaces by sampling bounded
+  actions from model-predicted Gaussian parameters.
+
+`Runner` interacts with Gymnasium environments and returns transition batches,
+episode returns, episode lengths, and optionally the last completed episode.
+
+`ReplayBuffer` stores transitions for off-policy learning. It supports both
+discrete and continuous action spaces through the configured action shape and
+action type.
+
+`TransitionBatch` is the tensor container passed between runners, replay
+buffers, experiments, and learners.
+
+## Simple Usage
+
+This example trains a DQN agent on `CartPole-v1`:
+
+```python
+import gymnasium as gym
+import torch as th
+
+from drlab import (
+    DQNConfig,
+    DQNLearner,
+    EpsilonGreedyController,
+    GreedyController,
+    OffPolicyExperiment,
+    OffPolicyExperimentConfig,
+)
+
+env = gym.make("CartPole-v1")
+
+obs_dim = env.observation_space.shape[0]
+num_actions = env.action_space.n
+
+model = th.nn.Sequential(
+    th.nn.Linear(obs_dim, 64),
+    th.nn.ReLU(),
+    th.nn.Linear(64, num_actions),
+)
+optimizer = th.optim.Adam(model.parameters(), lr=1e-3)
+
+learner = DQNLearner(
+    model,
+    optimizer,
+    DQNConfig(num_actions=num_actions),
+)
+
+controller = EpsilonGreedyController(
+    GreedyController(model, num_actions=num_actions),
+    num_actions=num_actions,
+    max_eps=1.0,
+    min_eps=0.05,
+    anneal_steps=10_000,
+)
+
+experiment = OffPolicyExperiment(
+    env,
+    controller,
+    learner,
+    OffPolicyExperimentConfig(
+        max_steps=10_000,
+        run_steps=1,
+        batch_size=64,
+        log_dir="runs/examples/dqn_cartpole",
+    ),
+)
+
+experiment.run()
+env.close()
+```
+
+More complete scripts are available in the [`examples/`](examples/) folder,
+including DQN, REINFORCE, actor-critic, PPO, SAC, and a comparison with
+Stable-Baselines3.

drlab-0.2.2/README.md

@@ -0,0 +1,171 @@
+# drlab
+
+`drlab` is a small deep reinforcement learning library for training neural
+network agents in Gymnasium environments. Its purpose is to provide simple,
+research-friendly building blocks for reinforcement learning experiments,
+especially when the environment or training setup is still being explored.
+
+The library is built around PyTorch models, Gymnasium environments, reusable
+controllers, replay buffers, runners, learners, and lightweight experiment
+loops. It is intended for research code and prototypes, not as a large
+production RL framework.
+
+## Installation
+
+The intended development setup uses [`uv`](https://docs.astral.sh/uv/). From the
+repository root:
+
+```bash
+uv sync --extra experiments --extra dev
+```
+
+Run commands through the managed environment:
+
+```bash
+uv run python -m pytest
+uv run python examples/dqn_cartpole.py
+```
+
+You can also activate the local virtual environment created by `uv`:
+
+```bash
+source .venv/bin/activate
+python -m pytest
+```
+
+If you prefer a normal `venv` and `pip` workflow:
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+python -m pip install --upgrade pip
+python -m pip install -e .
+```
+
+For experiment and development dependencies:
+
+```bash
+python -m pip install -e ".[experiments,dev]"
+```
+
+`drlab` depends on PyTorch, NumPy, Gymnasium, tqdm, and TensorBoard. If you need
+a specific CUDA build of PyTorch, install the matching PyTorch wheel before
+installing the package dependencies.
+
+## Library Overview
+
+Most public classes can be imported directly from `drlab`:
+
+```python
+from drlab import DQNConfig, DQNLearner, OffPolicyExperiment
+```
+
+### Off-Policy Learning
+
+Off-policy learners train from transition batches that may come from a replay
+buffer instead of only the most recent rollout. `OffPolicyExperiment` collects
+environment transitions with a `Runner`, stores them in a `ReplayBuffer`, samples
+minibatches, and calls the learner update.
+
+Implemented off-policy learners:
+
+- `DQNLearner`: value-based learning for discrete action spaces. It supports
+  replay-buffer training, target networks, Double DQN, hard or soft target
+  updates, gradient clipping, and custom regularizers.
+- `SACLearner`: Soft Actor-Critic style learning for continuous action spaces.
+  It uses an actor, two critics, target critics, entropy tuning, and replay
+  buffer training.
+
+### On-Policy Learning
+
+On-policy learners train from the data collected by the current policy.
+`OnPolicyExperiment` collects rollouts, computes returns when needed, and trains
+the learner directly on the fresh batch.
+
+Implemented on-policy learners:
+
+- `ReinforceLearner`: policy-gradient learning from discounted returns.
+- `ActorCriticLearner`: policy and value learning from shared model outputs.
+- `PPOLearner`: clipped policy optimization over rollout batches.
+
+### Supporting Components
+
+Controllers turn model outputs into environment actions:
+
+- `GreedyController`, `EpsilonGreedyController`, and `StochasticController`
+  support discrete action spaces.
+- `GaussianController` supports continuous action spaces by sampling bounded
+  actions from model-predicted Gaussian parameters.
+
+`Runner` interacts with Gymnasium environments and returns transition batches,
+episode returns, episode lengths, and optionally the last completed episode.
+
+`ReplayBuffer` stores transitions for off-policy learning. It supports both
+discrete and continuous action spaces through the configured action shape and
+action type.
+
+`TransitionBatch` is the tensor container passed between runners, replay
+buffers, experiments, and learners.
+
+## Simple Usage
+
+This example trains a DQN agent on `CartPole-v1`:
+
+```python
+import gymnasium as gym
+import torch as th
+
+from drlab import (
+    DQNConfig,
+    DQNLearner,
+    EpsilonGreedyController,
+    GreedyController,
+    OffPolicyExperiment,
+    OffPolicyExperimentConfig,
+)
+
+env = gym.make("CartPole-v1")
+
+obs_dim = env.observation_space.shape[0]
+num_actions = env.action_space.n
+
+model = th.nn.Sequential(
+    th.nn.Linear(obs_dim, 64),
+    th.nn.ReLU(),
+    th.nn.Linear(64, num_actions),
+)
+optimizer = th.optim.Adam(model.parameters(), lr=1e-3)
+
+learner = DQNLearner(
+    model,
+    optimizer,
+    DQNConfig(num_actions=num_actions),
+)
+
+controller = EpsilonGreedyController(
+    GreedyController(model, num_actions=num_actions),
+    num_actions=num_actions,
+    max_eps=1.0,
+    min_eps=0.05,
+    anneal_steps=10_000,
+)
+
+experiment = OffPolicyExperiment(
+    env,
+    controller,
+    learner,
+    OffPolicyExperimentConfig(
+        max_steps=10_000,
+        run_steps=1,
+        batch_size=64,
+        log_dir="runs/examples/dqn_cartpole",
+    ),
+)
+
+experiment.run()
+env.close()
+```
+
+More complete scripts are available in the [`examples/`](examples/) folder,
+including DQN, REINFORCE, actor-critic, PPO, SAC, and a comparison with
+Stable-Baselines3.

{drlab-0.1.1 → drlab-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "drlab"
-version = "0.1.1"
+version = "0.2.2"
 description = "Deep Reinforcement Learning kit for research."
 readme = "README.md"
 requires-python = ">=3.11"
@@ -20,7 +20,7 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
 ]
 dependencies = [
-    "torch==2.10.0+cu128",
+    "torch>=2.0,<3",
     "numpy==2.4.2",
     "gymnasium==1.2.3",
     "tqdm>=4.66",

drlab-0.2.2/src/drlab/__init__.py

@@ -0,0 +1,73 @@
+from drlab.controllers import (
+    ContinuousActionController,
+    Controller,
+    DiscreteActionController,
+    EpsilonGreedyController,
+    GaussianController,
+    GreedyController,
+    StochasticController,
+    WarmupController,
+)
+from drlab.experiments import (
+    OnPolicyExperiment,
+    OnPolicyExperimentConfig,
+    OffPolicyExperiment,
+    OffPolicyExperimentConfig,
+)
+from drlab.learners import (
+    ActorCriticConfig,
+    ActorCriticLearner,
+    DQNConfig,
+    DQNLearner,
+    SACConfig,
+    SACLearner,
+    OnPolicyConfig,
+    OnPolicyLearner,
+    OffPolicyConfig,
+    OffPolicyLearner,
+    PPOConfig,
+    PPOLearner,
+    ReinforceConfig,
+    ReinforceLearner,
+    TargetUpdate,
+    ValueTargets,
+)
+from drlab.replay import ReplayBuffer, TransitionBatch
+from drlab.runners import Runner
+
+__version__ = "0.2.2"
+
+__all__ = [
+    "__version__",
+    "ActorCriticConfig",
+    "ActorCriticLearner",
+    "OnPolicyExperiment",
+    "OnPolicyExperimentConfig",
+    "ContinuousActionController",
+    "Controller",
+    "DiscreteActionController",
+    "DQNConfig",
+    "DQNLearner",
+    "SACConfig",
+    "SACLearner",
+    "OnPolicyConfig",
+    "OnPolicyLearner",
+    "OffPolicyConfig",
+    "OffPolicyLearner",
+    "OffPolicyExperiment",
+    "OffPolicyExperimentConfig",
+    "PPOConfig",
+    "PPOLearner",
+    "ReinforceConfig",
+    "ReinforceLearner",
+    "TargetUpdate",
+    "EpsilonGreedyController",
+    "GaussianController",
+    "GreedyController",
+    "ReplayBuffer",
+    "Runner",
+    "StochasticController",
+    "TransitionBatch",
+    "ValueTargets",
+    "WarmupController",
+]

drlab-0.2.2/src/drlab/controllers/__init__.py

@@ -0,0 +1,17 @@
+from .base import ContinuousActionController, Controller, DiscreteActionController
+from .greedy import GreedyController
+from .e_greedy import EpsilonGreedyController
+from .stochastic import StochasticController
+from .gaussian import GaussianController
+from .warmup import WarmupController
+
+__all__ = [
+    "ContinuousActionController",
+    "Controller",
+    "DiscreteActionController",
+    "EpsilonGreedyController",
+    "GaussianController",
+    "GreedyController",
+    "StochasticController",
+    "WarmupController",
+]

drlab-0.2.2/src/drlab/controllers/base.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import torch as th
+from abc import ABC, abstractmethod
+
+
+class Controller(ABC):
+    """Base controller interface for all action spaces."""
+
+    model: th.nn.Module
+
+    @abstractmethod
+    def choose(self, obs: th.Tensor, **kwargs) -> th.Tensor: ...
+
+
+class DiscreteActionController(Controller):
+    """Controller interface for discrete action spaces."""
+
+    num_actions: int
+
+    @abstractmethod
+    def probabilities(self, obs: th.Tensor, **kwargs) -> th.Tensor: ...
+
+
+class ContinuousActionController(Controller):
+    """Controller interface for continuous action spaces."""
+
+    action_dim: int

{drlab-0.1.1 → drlab-0.2.2}/src/drlab/controllers/e_greedy.py

@@ -1,12 +1,13 @@
 import numpy as np
 import torch as th
-from .base import Controller
+from .base import DiscreteActionController
 
-class EpsilonGreedyController(Controller):
+
+class EpsilonGreedyController(DiscreteActionController):
 
     def __init__(
         self,
-        controller: Controller,
+        controller: DiscreteActionController,
         num_actions: int,
         max_eps: float = 1.0,
         min_eps: float = 0.1,
@@ -27,7 +28,12 @@ class EpsilonGreedyController(Controller):
         frac = max(1 - self.num_decisions / (self.anneal_steps - 1), 0.0)
         return frac * (self.max_eps - self.min_eps) + self.min_eps
 
-    def choose(self, obs: th.Tensor, increase_counter: bool = True, **kwargs) -> th.Tensor:
+    def choose(
+        self,
+        obs: th.Tensor,
+        increase_counter: bool = True,
+        **kwargs,
+    ) -> th.Tensor:
         eps = self.epsilon()
         if increase_counter:
             self.num_decisions += 1
@@ -40,7 +46,11 @@ class EpsilonGreedyController(Controller):
 
     def probabilities(self, obs: th.Tensor, **kwargs) -> th.Tensor:
         eps = self.epsilon()
-        greedy = self.controller.probabilities(obs, **kwargs)  # one-hot on argmax, shape [B,A]
+        greedy = self.controller.probabilities(obs, **kwargs)
         B = greedy.shape[0]
-        uniform = th.full((B, self.num_actions), eps / self.num_actions, device=greedy.device)
+        uniform = th.full(
+            (B, self.num_actions),
+            eps / self.num_actions,
+            device=greedy.device,
+        )
         return uniform + (1 - eps) * greedy

drlab-0.2.2/src/drlab/controllers/gaussian.py

@@ -0,0 +1,32 @@
+import torch as th
+from .base import ContinuousActionController
+
+
+class GaussianController(ContinuousActionController):
+
+    def __init__(
+        self,
+        model: th.nn.Module,
+        action_dim: int,
+        deterministic: bool = False,
+    ):
+        self.model = model
+        self.action_dim = action_dim
+        self.num_actions = action_dim
+        self.deterministic = deterministic
+
+    def choose(self, obs: th.Tensor, **kwargs) -> th.Tensor:
+        output = self.model(obs)
+
+        mean = output[:, :self.action_dim]
+        log_std = output[:, self.action_dim: 2 * self.action_dim]
+        log_std = th.clamp(log_std, -20, 2)
+
+        if self.deterministic:
+            u = mean
+        else:
+            std = log_std.exp()
+            eps = th.randn_like(mean)
+            u = mean + std * eps
+
+        return th.tanh(u)

{drlab-0.1.1 → drlab-0.2.2}/src/drlab/controllers/greedy.py

@@ -1,20 +1,20 @@
 import torch as th
-from .base import Controller
+from .base import DiscreteActionController
 
-class GreedyController(Controller):
+
+class GreedyController(DiscreteActionController):
 
     def __init__(self, model: th.nn.Module, num_actions: int):
-        
         self.num_actions = num_actions
         self.model = model
 
     def choose(self, obs: th.Tensor):
         output: th.Tensor = self.model(obs)[:, :self.num_actions]
         return th.argmax(output, dim=-1)
-    
+
     def probabilities(self, obs: th.Tensor):
         output: th.Tensor = self.model(obs)[:, :self.num_actions]
-        idx = output.argmax(dim=-1, keepdim=True)          # [B,1]
+        idx = output.argmax(dim=-1, keepdim=True)
         probs = th.zeros(output.shape[0], self.num_actions, device=output.device)
         probs.scatter_(dim=-1, index=idx, value=1.0)
         return probs

drlab-0.1.1/src/drlab/controllers/stochastic_controller.py → drlab-0.2.2/src/drlab/controllers/stochastic.py

@@ -1,17 +1,17 @@
 import torch as th
-from .base import Controller
+from .base import DiscreteActionController
 
-class StochasticController(Controller):
+
+class StochasticController(DiscreteActionController):
 
     def __init__(self, model: th.nn.Module, num_actions: int):
-        
         self.num_actions = num_actions
         self.model = model
 
     def choose(self, obs: th.Tensor) -> th.Tensor:
         probs = self.probabilities(obs)
         return th.distributions.Categorical(probs=probs).sample()
-        
+
     def probabilities(self, obs: th.Tensor) -> th.Tensor:
         output = self.model(obs)[:, :self.num_actions]
-        return th.nn.functional.softmax(output, dim=-1)
+        return th.nn.functional.softmax(output, dim=-1)

drlab-0.2.2/src/drlab/controllers/warmup.py

@@ -0,0 +1,42 @@
+from typing import Any
+
+import numpy as np
+import torch as th
+
+from .base import ContinuousActionController, Controller, DiscreteActionController
+
+
+class WarmupController(ContinuousActionController, DiscreteActionController):
+    def __init__(self, controller: Controller, action_space: Any, warmup_steps: int):
+        self.controller = controller
+        self.action_space = action_space
+        self.warmup_steps = warmup_steps
+        self.steps = 0
+        self.model = controller.model
+
+        if hasattr(controller, "num_actions"):
+            self.num_actions = controller.num_actions
+        if hasattr(controller, "action_dim"):
+            self.action_dim = controller.action_dim
+
+        if warmup_steps < 0:
+            raise ValueError("warmup_steps must be >= 0.")
+
+    def choose(
+        self,
+        obs: th.Tensor,
+        **kwargs,
+    ) -> th.Tensor:
+        if self.steps >= self.warmup_steps:
+            return self.controller.choose(obs, **kwargs)
+
+        batch_size = obs.shape[0]
+        self.steps += batch_size
+        actions = [self.action_space.sample() for _ in range(batch_size)]
+        return th.as_tensor(np.asarray(actions), device=obs.device)
+
+    def probabilities(self, obs: th.Tensor, **kwargs) -> th.Tensor:
+        if self.steps >= self.warmup_steps:
+            return self.controller.probabilities(obs, **kwargs)
+
+        return th.ones((obs.shape[0], self.num_actions), device=obs.device) / self.num_actions

drlab-0.2.2/src/drlab/experiments/__init__.py

@@ -0,0 +1,9 @@
+from .off_policy_experiment import OffPolicyExperiment, OffPolicyExperimentConfig
+from .on_policy_experiment import OnPolicyExperiment, OnPolicyExperimentConfig
+
+__all__ = [
+    "OffPolicyExperiment", 
+    "OnPolicyExperiment",
+    "OffPolicyExperimentConfig", 
+    "OnPolicyExperimentConfig",
+]