drlab 0.1.0__py3-none-any.whl

drlab/__init__.py

@@ -0,0 +1,36 @@
+from drlab.controllers import (
+    Controller,
+    EpsilonGreedyController,
+    GreedyController,
+    StochasticController,
+)
+from drlab.experiments import (
+    ActorCriticExperiment,
+    ActorCriticExperimentConfig,
+    DQNExperiment,
+    DQNExperimentConfig,
+)
+from drlab.learners import ActorCritic, ActorCriticConfig, DQN, DQNConfig
+from drlab.replay import ReplayBuffer, TransitionBatch
+from drlab.runners import Runner
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "__version__",
+    "ActorCritic",
+    "ActorCriticConfig",
+    "ActorCriticExperiment",
+    "ActorCriticExperimentConfig",
+    "Controller",
+    "DQN",
+    "DQNConfig",
+    "DQNExperiment",
+    "DQNExperimentConfig",
+    "EpsilonGreedyController",
+    "GreedyController",
+    "ReplayBuffer",
+    "Runner",
+    "StochasticController",
+    "TransitionBatch",
+]

drlab/controllers/__init__.py

@@ -0,0 +1,6 @@
+from .base import Controller
+from .greedy import GreedyController
+from .e_greedy import EpsilonGreedyController
+from .stochastic_controller import StochasticController
+
+__all__ = ["Controller", "GreedyController", "EpsilonGreedyController", "StochasticController"]

drlab/controllers/base.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+import torch as th
+from abc import ABC, abstractmethod
+
+class Controller(ABC):
+    """Abstract controller interface."""
+    
+    num_actions: int
+    model: th.nn.Module = None  
+    controller: Controller = None
+
+    @abstractmethod
+    def choose(self, obs: th.Tensor, **kwargs) -> th.Tensor: ...
+
+    @abstractmethod
+    def probabilities(self, obs: th.Tensor, **kwargs) -> th.Tensor: ...

drlab/controllers/e_greedy.py

@@ -0,0 +1,46 @@
+import numpy as np
+import torch as th
+from .base import Controller
+
+class EpsilonGreedyController(Controller):
+
+    def __init__(
+        self,
+        controller: Controller,
+        num_actions: int,
+        max_eps: float = 1.0,
+        min_eps: float = 0.1,
+        anneal_steps: int = 10_000,
+    ):
+        self.controller = controller
+        self.num_actions = num_actions
+        self.model = controller.model
+        self.max_eps = max_eps
+        self.min_eps = min_eps
+        self.anneal_steps = anneal_steps
+        self.num_decisions = 0
+
+        if anneal_steps <= 1:
+            raise ValueError("anneal_steps must be >= 2")
+
+    def epsilon(self) -> float:
+        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:
+        eps = self.epsilon()
+        if increase_counter:
+            self.num_decisions += 1
+
+        B = obs.shape[0] if obs.ndim > 1 else 1
+        if np.random.rand() < eps:
+            return th.randint(self.num_actions, (B,), device=obs.device, dtype=th.long)
+
+        return self.controller.choose(obs, **kwargs)
+
+    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]
+        B = greedy.shape[0]
+        uniform = th.full((B, self.num_actions), eps / self.num_actions, device=greedy.device)
+        return uniform + (1 - eps) * greedy

drlab/controllers/greedy.py

@@ -0,0 +1,20 @@
+import torch as th
+from .base import Controller
+
+class GreedyController(Controller):
+
+    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]
+        probs = th.zeros(output.shape[0], self.num_actions, device=output.device)
+        probs.scatter_(dim=-1, index=idx, value=1.0)
+        return probs

drlab/controllers/stochastic_controller.py

@@ -0,0 +1,17 @@
+import torch as th
+from .base import Controller
+
+class StochasticController(Controller):
+
+    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)

drlab/experiments/__init__.py

@@ -0,0 +1,9 @@
+from .dqn_experiment import DQNExperiment, DQNExperimentConfig
+from .ac_experiment import ActorCriticExperiment, ActorCriticExperimentConfig
+
+__all__ = [
+    "DQNExperimentConfig", 
+    "DQNExperiment", 
+    "ActorCriticExperimentConfig", 
+    "ActorCriticExperiment"
+]

drlab/experiments/ac_experiment.py

@@ -0,0 +1,102 @@
+import numpy as np
+from tqdm import tqdm
+import gymnasium as gym
+from dataclasses import dataclass
+from typing import Callable
+from torch.utils.tensorboard import SummaryWriter
+
+from drlab.learners import ActorCritic
+from drlab.runners import Runner
+from drlab.controllers import Controller
+
+
+@dataclass
+class ActorCriticExperimentConfig:
+    max_steps: int
+    gamma: float = 0.99
+    run_steps: int = 0
+    log_dir: str = "runs/reinforce_experiment"
+    experiment_name: str = "ActorCriticExperiment"
+    step_callback: Callable[[int], None] | None = None
+    step_callback_interval: int | None = None
+
+
+class ActorCriticExperiment:
+
+    def __init__(
+        self,
+        env: gym.Env,
+        controller: Controller,
+        learner: ActorCritic,
+        config: ActorCriticExperimentConfig,
+    ):  
+        # Init experiment settings
+        self.max_steps = config.max_steps
+        self.run_steps = config.run_steps
+        self.step_callback = config.step_callback
+        self.step_callback_interval = config.step_callback_interval
+        
+        # Only ask the runner to compute discounted returns when the learner uses them.
+        calculate_returns = (
+            not learner.config.advantage_bootstrap
+            or (learner.config.use_bias and learner.config.value_targets == "returns")
+        )
+
+        # Init drl components
+        self.runner = Runner(env, controller, calculate_returns, False, config.gamma, learner.device)
+        self.learner = learner
+
+        # Init logging
+        self.writer = SummaryWriter(log_dir=config.log_dir)
+        self.experiment_name = config.experiment_name
+
+        if self.step_callback is not None and self.step_callback_interval is None:
+            raise ValueError("step_callback_interval must be set when step_callback is provided.")
+
+    def run(self):
+
+        steps = 0
+        if self.step_callback is not None:
+            self.step_callback(steps)
+        next_callback_step = self.step_callback_interval
+        pbar = tqdm(total=self.max_steps, desc=self.experiment_name, dynamic_ncols=True, mininterval=1.0)
+        while steps < self.max_steps:
+            
+            # 1. Run batch of transitions
+            batch, ep_returns, ep_lengths, _ = self.runner.run(self.run_steps)
+            batch = batch.to(self.learner.device)
+
+            # 2. Learn from batch
+            loss = self.learner.train(
+                batch.rewards,
+                batch.dones,
+                batch.states,
+                batch.actions,
+                batch.next_states,
+                batch.returns,
+            )
+
+            # 3. Log results
+            self.writer.add_scalar("Loss", loss, steps)
+            if ep_returns:
+                mean_return = np.mean(ep_returns)
+                mean_length = np.mean(ep_lengths)
+                self.writer.add_scalar("Episode return", mean_return, steps)
+                self.writer.add_scalar("Episode Length", mean_length, steps)
+
+                pbar.set_postfix({
+                    "loss": f"{loss:.3f}",
+                    "return": f"{mean_return:.2f}",
+                    "len": f"{mean_length:.1f}",
+                }, refresh=False)
+
+            step_inc = batch.states.shape[0]
+            steps += step_inc
+            pbar.update(step_inc)
+
+            while self.step_callback is not None and steps >= next_callback_step:
+                self.step_callback(next_callback_step)
+                next_callback_step += self.step_callback_interval
+
+        pbar.close()
+        self.writer.close()

drlab/experiments/dqn_experiment.py

@@ -0,0 +1,151 @@
+import numpy as np
+from tqdm import tqdm
+import gymnasium as gym
+from dataclasses import dataclass
+from typing import Callable
+from torch.utils.tensorboard import SummaryWriter
+
+from drlab.learners import DQN
+from drlab.runners import Runner
+from drlab.controllers import Controller
+from drlab.replay import TransitionBatch, ReplayBuffer
+
+@dataclass
+class DQNExperimentConfig:
+    max_steps: int
+    gamma: float = 0.99
+    run_steps: int = 0
+    log_dir: str = "runs/dqn_experiment"
+    experiment_name: str = "DQNExperiment"
+    use_replay: bool = True
+    replay_buffer_size: int = 10_000
+    batch_size: int = 128
+    use_last_episode: bool = True
+    grad_repeats: int = 1
+    step_callback: Callable[[int], None] | None = None
+    step_callback_interval: int | None = None
+
+
+class DQNExperiment:
+
+    def __init__(
+            self,
+            env: gym.Env,
+            controller: Controller,
+            learner: DQN,
+            config: DQNExperimentConfig,
+        ):
+
+        # Init experiment settings
+        self.max_steps = config.max_steps
+        self.grad_repeats = config.grad_repeats
+        self.batch_size = config.batch_size
+        self.use_last_episode = config.use_last_episode
+        self.run_steps = config.run_steps
+        self.step_callback = config.step_callback
+        self.step_callback_interval = config.step_callback_interval
+
+        # Init drl components
+        self.runner = Runner(env, controller, False, True, config.gamma, learner.device)
+        self.learner = learner
+
+        # Init replay buffer
+        self.replay_buffer_size = config.replay_buffer_size
+        self.use_replay = config.use_replay
+        self.replay_buffer = ReplayBuffer(
+            capacity=config.replay_buffer_size if config.use_replay else config.batch_size,  
+            obs_shape=env.observation_space.shape,
+            device=learner.device
+        )
+
+        # Init logging
+        self.writer = SummaryWriter(log_dir=config.log_dir)
+        self.experiment_name = config.experiment_name
+
+        if self.step_callback is not None and self.step_callback_interval is None:
+            raise ValueError("step_callback_interval must be set when step_callback is provided.")
+
+    def _make_minibatch(self, batch: TransitionBatch, last_episode: TransitionBatch | None) -> TransitionBatch:
+
+        # No replay buffer
+        if not self.use_replay:
+            return self.replay_buffer.get_all()
+        
+        # Replay buffer without last episode
+        if not self.use_last_episode:
+            return self.replay_buffer.sample(self.batch_size)
+
+        # Replay buffer with last episode
+        episode = last_episode if last_episode else batch
+        episode = episode.to(self.learner.device)
+
+        ep_len = episode.states.shape[0]
+        if ep_len >= self.batch_size:
+            return episode
+
+        rest = self.replay_buffer.sample(self.batch_size - ep_len)
+        return rest.cat(episode)
+
+    def _learn_from_batch(self, batch: TransitionBatch, last_episode: TransitionBatch | None) -> float:
+        # 1) Add batch to replay buffer
+        self.replay_buffer.add(
+            states=batch.states.cpu().numpy(),
+            actions=batch.actions.cpu().numpy(),
+            rewards=batch.rewards.cpu().numpy(),
+            dones=batch.dones.cpu().numpy(),
+            next_states=batch.next_states.cpu().numpy(),
+            returns=batch.returns.cpu().numpy(),
+        )
+        if self.replay_buffer.size < self.batch_size:
+            return 0.0
+
+        # 2) train repeats
+        total_loss = 0.0
+        for _ in range(self.grad_repeats):
+            mb = self._make_minibatch(batch, last_episode)
+            total_loss += self.learner.train(
+                mb.rewards, mb.dones, mb.states, mb.actions, mb.next_states
+            )
+
+        return total_loss / self.grad_repeats
+
+    def run(self):
+
+        steps = 0
+        if self.step_callback is not None:
+            self.step_callback(steps)
+        next_callback_step = self.step_callback_interval
+        pbar = tqdm(total=self.max_steps, desc=self.experiment_name, dynamic_ncols=True, mininterval=1.0)
+        while steps < self.max_steps:
+
+            # 1. Run batch of transitions
+            batch, ep_returns, ep_lengths, last_episode = self.runner.run(self.run_steps)
+
+            # 2. Learn from batch
+            loss = self._learn_from_batch(batch, last_episode)
+
+            # 3. Log results
+            self.writer.add_scalar("Loss", loss, steps)
+            if ep_returns:
+                mean_return = np.mean(ep_returns)
+                mean_length = np.mean(ep_lengths)
+                self.writer.add_scalar("Episode return", mean_return, steps)
+                self.writer.add_scalar("Episode Length", mean_length, steps)
+
+                pbar.set_postfix({
+                    "loss": f"{loss:.3f}",
+                    "return": f"{mean_return:.2f}",
+                    "len": f"{mean_length:.1f}",
+                }, refresh=False)
+
+            # 3. Update steps & progress bar
+            step_inc = batch.states.shape[0] 
+            steps += step_inc 
+            pbar.update(step_inc)
+
+            while self.step_callback is not None and steps >= next_callback_step:
+                self.step_callback(next_callback_step)
+                next_callback_step += self.step_callback_interval
+        
+        pbar.close()
+        self.writer.close()

drlab/learners/__init__.py

@@ -0,0 +1,4 @@
+from .dqn import DQN, DQNConfig
+from .actor_critic import ActorCritic, ActorCriticConfig
+
+__all__ = ["DQNConfig", "DQN", "ActorCriticConfig", "ActorCritic"]

drlab/learners/actor_critic.py

@@ -0,0 +1,206 @@
+from dataclasses import dataclass, field
+from typing import Callable
+
+import torch as th
+
+@dataclass
+class ActorCriticConfig:
+    device: th.device | str = "cpu"
+    regularizers: list[Callable[..., th.Tensor | float]] = field(default_factory=list)
+    reg_lams: list[float] = field(default_factory=list)
+    num_actions: int = 2
+    clip_grad: bool = True
+    grad_norm_clip: float = 1.0 
+    use_bias: bool = True
+    value_criterion: Callable = field(default_factory=th.nn.MSELoss)
+    value_lambda: float = 0.1
+    value_targets: str = "td"
+    gamma: float = 0.99
+    advantage_bootstrap: bool = True
+    off_policy_iterations: int = 0
+    ppo_clipping: float = 0.1
+    use_entropy: bool = False
+    entropy_max_lambda: float = 0.0
+    entropy_min_lambda: float = 0.0
+    entropy_anneal_steps: int = 1
+    normalize_advantages: bool = False
+
+
+class ActorCritic:
+
+    def __init__(
+        self,
+        actor: th.nn.Module,
+        optimizer: th.optim.Optimizer,
+        config: ActorCriticConfig,
+       ):
+        
+        device = th.device(config.device)
+        self.actor = actor
+        self.optimizer = optimizer
+        self.all_parameters = list(actor.parameters())
+        self.actor.to(device)
+        self.device = device
+        self.config = config
+
+        if len(config.regularizers) != len(config.reg_lams):
+            raise ValueError("Length of regularizers and reg_lams lists must match.")
+        
+        if not config.use_bias and config.advantage_bootstrap:
+            raise ValueError("advantage_bootstrap=True requires use_bias=True.")
+        
+        if config.value_targets not in {"returns", "td"}:
+            raise ValueError("value_targets must be either 'returns' or 'td'.")
+
+        if config.use_entropy and config.entropy_anneal_steps <= 1:
+            raise ValueError("entropy_anneal_steps must be > 1 when use_entropy=True.")
+        if config.use_entropy and config.entropy_max_lambda < config.entropy_min_lambda:
+            raise ValueError("entropy_max_lambda must be >= entropy_min_lambda.")
+
+        self.entropy_step = 0
+        
+    def _value_loss(
+        self,
+        returns: th.Tensor,
+        rewards: th.Tensor,
+        dones: th.Tensor,
+        values: th.Tensor,
+        next_values: th.Tensor,
+    ) -> th.Tensor:
+        if self.config.value_targets == "returns":
+            targets = returns
+        elif self.config.value_targets == "td":
+            targets = rewards + self.config.gamma * (~dones * next_values)
+        else:
+            raise ValueError(f"Unknown value_targets: {self.config.value_targets}")
+
+        return self.config.value_criterion(values, targets)
+    
+    def _get_policy(self, logits: th.Tensor, actions: th.Tensor) -> th.Tensor:
+        probs = th.nn.functional.softmax(logits, dim=-1)
+        pi = probs.gather(dim=-1, index=actions)
+        return pi
+
+    def _entropy_lambda(self) -> float:
+        progress = min(self.entropy_step / (self.config.entropy_anneal_steps - 1), 1.0)
+
+        return (
+            self.config.entropy_max_lambda
+            + progress * (self.config.entropy_min_lambda - self.config.entropy_max_lambda)
+        )
+
+    def _entropy_loss(self, logits: th.Tensor) -> th.Tensor:
+        entropy_lambda = self._entropy_lambda()
+        probs = th.nn.functional.softmax(logits, dim=-1)
+        log_probs = th.nn.functional.log_softmax(logits, dim=-1)
+        entropy = -(probs * log_probs).sum(dim=-1).mean()
+        return -entropy_lambda * entropy
+        
+    def _advantages(
+            self, 
+            returns: th.Tensor, 
+            rewards: th.Tensor,
+            dones: th.Tensor,
+            values: th.Tensor,
+            next_values: th.Tensor
+        ) -> th.Tensor:
+        advantages = None
+        if self.config.advantage_bootstrap:
+            advantages = rewards + self.config.gamma * (~dones * next_values)
+        else: 
+            advantages = returns
+        if self.config.use_bias:
+            advantages = advantages - values.detach()
+
+        if self.config.normalize_advantages and advantages.numel() > 1:
+            advantages = (advantages - advantages.mean()) / (
+                advantages.std(unbiased=False) + 1e-8
+            )
+
+        return advantages
+
+    def _policy_loss(self, pi: th.Tensor, advantage: th.Tensor) -> th.Tensor:
+
+        if self.old_pi is None:
+            self.old_pi = pi.detach()
+            return -th.mean(pi.clamp_min(1e-8).log() * advantage.detach())
+        else:
+            ratios = pi / self.old_pi.detach()
+            loss = advantage.detach() * ratios
+            ppo_loss = th.clamp(ratios, 1-self.config.ppo_clipping, 1+self.config.ppo_clipping) * advantage.detach()
+            loss = th.min(loss, ppo_loss)
+            return -th.mean(loss)
+
+    def train(
+        self,
+        rewards: th.Tensor,      # float32, [B,1]
+        dones: th.Tensor,        # bool or float(0/1), [B,1]
+        states: th.Tensor,       # float32, [B, obs_dim] or [B,C,H,W]
+        actions: th.Tensor,      # int64, [B,1]
+        next_states: th.Tensor,  # float32, same as states
+        returns: th.Tensor,      # float32, [B,1]
+    ) -> float:
+        
+        self.actor.train(True)
+        need_next_values = (
+            self.config.advantage_bootstrap
+            or self.config.value_targets == "td"
+        )
+
+        self.old_pi, loss_sum = None, 0
+
+        for _ in range(1 + self.config.off_policy_iterations):
+            # 1. Compute model policy, values and next_values
+            output: th.Tensor  = self.actor(states)
+            logits = output[:, :self.config.num_actions]
+            values = output[:, self.config.num_actions:self.config.num_actions + 1]
+
+            next_values = None
+            if need_next_values:
+                with th.no_grad():
+                    next_output: th.Tensor  = self.actor(next_states)
+                    next_values = next_output[:, self.config.num_actions:self.config.num_actions + 1]
+            pi = self._get_policy(logits, actions)
+
+            # 2. Compute losses
+            policy_loss = self._policy_loss(
+                pi, 
+                self._advantages(
+                    returns, 
+                    rewards,
+                    dones,
+                    values,
+                    next_values
+                )
+            )
+            entropy_loss = 0.0
+            if self.config.use_entropy:
+                entropy_loss = self._entropy_loss(logits)
+            value_loss = 0.0
+            if self.config.use_bias:
+                value_loss = self.config.value_lambda * self._value_loss(
+                    returns,
+                    rewards,
+                    dones,
+                    values,
+                    next_values
+                )
+            reg_loss = (
+                sum(
+                    lam * reg(self.actor, rewards, dones, states, actions, next_states)
+                    for reg, lam in zip(self.config.regularizers, self.config.reg_lams)
+                )
+                if self.config.regularizers else 0.0
+            )
+            loss = policy_loss + value_loss + entropy_loss + reg_loss
+
+            # 3. Optimize
+            self.optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            if self.config.clip_grad:
+                th.nn.utils.clip_grad_norm_(self.all_parameters, self.config.grad_norm_clip)
+            self.optimizer.step()
+            loss_sum += loss.item()
+
+        self.entropy_step += 1
+        return loss_sum