zeroRl 0.1.1__py3-none-any.whl

zerorl/algorithms/ppo/__init__.py

@@ -0,0 +1,9 @@
+"""Proximal Policy Optimization algorithm.
+
+Re-exports gae_compute, ppo_func, and easy_train_ppo.
+"""
+
+from .ppo import gae_compute, ppo_func
+from  .easy_ppo import easy_train_ppo
+
+__all__ = ["gae_compute",  "ppo_func", "easy_train_ppo"]

zerorl/algorithms/ppo/easy_ppo.py

@@ -0,0 +1,70 @@
+"""One-call PPO quickstart.
+
+Provides easy_train_ppo() which wires together agent, environment, buffer,
+and PPO update into a ready-to-train BaseTrain instance.
+"""
+
+from typing import Callable
+from torch import optim
+from zerorl.factory import get_env, get_actor_critic_buffer, ActorCriticAgent
+from zerorl.train import BaseTrain
+from zerorl.algorithms.ppo import ppo_func, gae_compute
+from zerorl.config import TrainConfig, AlgoConfig
+from zerorl.helpers.agent import BaseAgent
+from zerorl.helpers.env import BaseEnv
+from zerorl.functions import get_obs_act
+
+def easy_train_ppo(env_spec: str | Callable | BaseEnv, 
+                    config: TrainConfig,
+                    algo_config: AlgoConfig,
+                    hidden_layer: int = 64,
+                    render_mode: str | None = None,
+                    base_agent: BaseAgent | None = None,
+                    optimizer: optim.Optimizer | None = None,
+                    schedule_func: Callable[[int], float] | None = None,
+                   ):
+    """Create and return a BaseTrain instance with PPO wiring.
+
+    Automatically builds the environment, agent, buffer, and update function.
+
+    Args:
+        env_spec: Gymnasium env ID string, BaseEnv subclass/instance, or callable.
+        config: Training configuration (device, paths, timesteps).
+        algo_config: Algorithm hyperparameters (lr, gamma, clip_eps, etc.).
+        hidden_layer: Hidden layer size for the default ActorCriticAgent.
+        render_mode: Render mode for the environment.
+        base_agent: Custom agent (overrides the default ActorCriticAgent).
+        optimizer: Custom optimizer (overrides default Adam).
+        schedule_func: Custom LR schedule function (overrides default linear decay).
+
+    Returns:
+        BaseTrain instance ready for .train() or .test().
+    """
+    env = get_env(env_spec, config.num_envs, render_mode)
+    obs_dim, act_dim, n_obs, n_act, is_discrete = get_obs_act(env)
+    
+    if base_agent is not None:
+        agent = base_agent
+    else:
+        agent = ActorCriticAgent(n_obs, n_act, is_discrete, hidden_layer) #type: ignore
+
+    buffer = get_actor_critic_buffer(obs_dim, act_dim, config) #type: ignore
+
+    #update weights function
+    def easy_update_weights(agent, buffer, scheduler, optimizer, last_output, algo_config):
+        data = buffer.get_all()
+        gae_compute(data["reward"], data["value"], last_output["value"], data["done"], buffer, algo_config)
+        return ppo_func(agent, optimizer, buffer, algo_config, scheduler, device=agent.device)
+
+    train = BaseTrain(
+            agent = agent,
+            env = env,
+            buffer = buffer,
+            update_weights = easy_update_weights,
+            config = config,
+            algo_config = algo_config,
+            optimizer = optimizer,
+            schedule_func = schedule_func,
+            render_mode = render_mode,
+            )
+    return train

zerorl/algorithms/ppo/ppo.py

@@ -0,0 +1,232 @@
+"""Proximal Policy Optimization (PPO) standalone functions.
+
+Provides gae_compute(), ppo_loss(), and ppo_func() for computing GAE
+advantages, computing the clipped surrogate loss, and running the PPO
+optimization step.
+Reference: Schulman et al., "Proximal Policy Optimization Algorithms" (2017)
+"""
+
+import torch
+from typing import Callable
+from torch import nn
+from torch import Tensor
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LambdaLR
+from zerorl.buffer import Buffer
+from zerorl.helpers.agent import BaseAgent, eval_action
+from zerorl.config import AlgoConfig
+from zerorl.functions import get_buffer_params_model, fast_compile
+from zerorl.errors import assert_agent_contract
+
+
+def gae_compute(rewards: Tensor,
+            values: Tensor,
+            last_value: Tensor,
+            dones: Tensor,
+            buffer:Buffer,
+            algo_config: AlgoConfig):
+    """Compute Generalized Advantage Estimation.
+
+    Works backwards through the trajectory, accumulating TD errors
+    with exponentially decaying weights. Writes "advantage" and "return"
+    directly into buffer.data.
+
+    Args:
+        rewards: Rewards for each timestep, shape (T, num_envs).
+        values: Value estimates for each timestep, shape (T, num_envs).
+        last_value: Bootstrap value for the state after the last step, shape (num_envs,).
+        dones: Episode termination flags, shape (T, num_envs). 1.0 = done.
+        buffer: Buffer to write "advantage" and "return" into.
+        algo_config: Algorithm configuration (gamma, gae_lambda).
+    """
+    rewards = rewards.reshape(rewards.shape[0], -1)
+    values = values.reshape(values.shape[0], -1)
+    dones = dones.reshape(dones.shape[0], -1)
+    last_value = last_value.reshape(-1)
+    num_envs = rewards.shape[1]
+    gae = torch.zeros(num_envs, dtype=torch.float32, device=rewards.device)
+    # Mask: 0.0 at episode boundaries (no bootstrapping across episodes)
+    mask = 1.0 - dones
+    next_values = torch.cat((values[1:], last_value.unsqueeze(0)), 0)
+    total_size = rewards.shape[0]
+    delta = rewards + algo_config.gamma * next_values * mask - values
+    advantages = torch.empty_like(delta)
+    for step in reversed(range(total_size)):
+        gae = delta[step] + algo_config.gamma * algo_config.gae_lambda * mask[step] * gae
+        advantages[step] = gae
+    returns = advantages + values
+    buffer.data["advantage"][:buffer.size] = advantages
+    buffer.data["return"][:buffer.size] = returns
+
+
+def ppo_loss(
+        agent: BaseAgent,
+        params: dict,
+        buffers: dict,
+        states: Tensor,
+        actions: Tensor,
+        old_log_prob: Tensor,
+        old_values: Tensor,
+        advantages: Tensor,
+        returns: Tensor,
+        ent_coef: float,
+        value_coef: float,
+        clip_eps: float,
+        clip_vf: float,
+        ) -> dict[str, Tensor]:
+    """Compute PPO clipped surrogate loss, value loss, and entropy bonus.
+
+    Args:
+        agent: The policy network.
+        params: Named parameters dict from get_buffer_params_model().
+        buffers: Named buffers dict from get_buffer_params_model().
+        states: Batch of observations.
+        actions: Batch of actions taken.
+        old_log_prob: Log probabilities from the old policy.
+        old_values: Value estimates from the old policy.
+        advantages: GAE advantage estimates.
+        returns: GAE return estimates.
+        ent_coef: Entropy bonus coefficient.
+        value_coef: Value loss coefficient.
+        clip_eps: PPO clipping range.
+        clip_vf: Whether to clip value predictions.
+
+    Returns:
+        Dict with keys "loss", "policy_loss", "value_loss", "entropy_loss".
+    """
+    logits, new_values = torch.func.functional_call(agent, (params, buffers), (states,))
+    dist = agent.build_distribution(logits) #type: ignore[operator]
+    new_log_probs, dist_entropy = eval_action(dist, actions)
+
+    idx_adv = advantages.view(-1)
+    idx_return = returns.view(-1)
+    new_values = new_values.view(-1)
+    old_values = old_values.view(-1)
+    old_log_prob = old_log_prob.view(-1)
+
+    logratio = new_log_probs - old_log_prob
+    ratio = torch.exp(logratio)
+
+    clip_eps = clip_eps
+    surr1 = ratio * idx_adv
+    surr2 = torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps) * idx_adv
+    policy_loss = -torch.min(surr1, surr2).mean()
+    
+    clip_vf = clip_vf
+    if clip_vf:
+        value_pred_clipped = old_values + (new_values - old_values).clamp(-clip_eps, clip_eps)
+        value_loss = 0.5 * torch.max((idx_return - new_values).pow(2), (value_pred_clipped - idx_return).pow(2)).mean() 
+    else:
+        value_loss = 0.5 * nn.functional.mse_loss(new_values, idx_return)
+
+    entropy_loss = dist_entropy.mean()
+
+    loss = policy_loss + \
+            (value_coef * value_loss) - \
+            (ent_coef * entropy_loss)
+    return {'loss': loss,
+            'policy_loss': policy_loss,
+            'value_loss': value_loss,
+            'entropy_loss':entropy_loss}
+
+
+def ppo_func(agent: BaseAgent,
+        optimizer: Optimizer,
+        buffer: Buffer,
+        algo_config: AlgoConfig,
+        scheduler: LambdaLR,
+        *,
+        ppo_loss_func: Callable[[BaseAgent, dict, dict, Tensor, Tensor,
+                                 Tensor, Tensor, Tensor, Tensor, float,
+                                 float, float, float], dict[str, Tensor]] = ppo_loss,
+        device: torch.device = torch.device("cpu")
+        ) ->  dict[str, Tensor]:
+    """Run a full PPO update on collected rollout data.
+
+    Steps the scheduler first, then normalizes advantages and runs
+    minibatch SGD epochs with clipped surrogate loss and gradient clipping.
+
+    Args:
+        agent: The policy network.
+        optimizer: Optimizer for the agent parameters.
+        buffer: Buffer containing rollout data with keys "state", "action",
+            "log_prob", "value", "advantage", "return".
+        algo_config: Algorithm configuration (batch_size, epochs, clip_eps, etc.).
+        scheduler: Learning rate scheduler (stepped once per call).
+        device: Torch device for computations.
+
+    Returns:
+        Dict of averaged loss metrics ("loss", "policy_loss", "value_loss", "entropy_loss").
+    """
+    assert_agent_contract(agent,
+         {"forward": "Your agent should have the method `forward`",
+          "build_distribution": "Your agent should have the method `build_distribution`"})
+
+    params, buffers = get_buffer_params_model(agent)
+    all_data = buffer.get_all()
+    flat_data = {key: tensor.reshape(-1, *tensor.shape[2:])
+                for key, tensor in all_data.items()}
+
+    mb_advantages = flat_data["advantage"]
+    adv_norm = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)
+    returns = flat_data["return"]
+
+    #coefficient and eps for clipping
+    value_coef = algo_config.value_coef
+    ent_coef = algo_config.ent_coef
+    clip_eps = algo_config.clip_eps
+    clip_vf = getattr(algo_config, "clip_vf", False)
+
+    dataset_size = flat_data["action"].size(0)
+    final_metrics: dict[str, Tensor] = {}
+
+    @fast_compile(mode="reduce-overhead") #type: ignore
+    def ppo_backward(agent: BaseAgent,
+                    params: dict,
+                    buffers: dict,
+                    state: Tensor,
+                    action: Tensor,
+                    old_log_prob: Tensor,
+                    old_values: Tensor,
+                    advantage: Tensor,
+                    return_: Tensor) -> dict[str, Tensor]:
+        global_losses = ppo_loss_func(agent, params, buffers, state, action,
+                                old_log_prob, old_values, advantage, return_,
+                                value_coef, ent_coef, clip_eps, clip_vf)
+        loss_tensor = global_losses["loss"]
+        loss_tensor.backward()
+        return global_losses
+
+    def update() -> list[dict[str, Tensor]]:
+        """Run multiple epochs of minibatch SGD on rollout data.
+
+        Shuffles data each epoch, splits into minibatches, and updates weights
+        with gradient clipping.
+        """
+        history = []
+        for _ in range(algo_config.epochs):
+            shuffle_index = torch.randperm(dataset_size, device=device)
+            for start in range(0, dataset_size, algo_config.batch_size):
+                end = start + algo_config.batch_size
+                idx = shuffle_index[start:end]
+                optimizer.zero_grad(set_to_none=True)
+                torch.compiler.cudagraph_mark_step_begin()
+                global_losses: dict[str, Tensor] = ppo_backward(agent, params, buffers, flat_data["state"][idx],
+                                flat_data["action"][idx], flat_data["log_prob"][idx],
+                                flat_data["value"][idx], adv_norm[idx], returns[idx])
+                torch.nn.utils.clip_grad_norm_(agent.parameters(), 0.5)
+                optimizer.step()
+                history.append({k: v.clone().detach() for k, v in global_losses.items()})
+        return history
+
+    # Compute losses and update weights
+    history = update()
+    scheduler.step()
+    # Return average losses over all actual updates ([:index_loss] excludes
+    # any unused pre-allocated entries)
+    if history:
+        keys = history[0].keys()
+        for key in keys:
+            stacked = torch.stack([h[key] for h in history])
+            final_metrics[key] = stacked.mean()
+    return final_metrics

zerorl/buffer.py

@@ -0,0 +1,80 @@
+"""Pre-allocated rollout buffer for RL training.
+
+Provides Buffer, which stores trajectory data in fixed-size PyTorch
+tensors and converts them for the PPO update step.
+"""
+
+import torch
+from zerorl.errors import KeyBufferError
+from zerorl.config import TrainConfig
+
+
+class Buffer:
+    """Pre-allocated rollout buffer for collecting RL trajectory data.
+
+    Stores trajectory data in pre-allocated PyTorch tensors with a slice
+    pointer for O(1) insertion. gae_compute() writes advantage and return
+    directly into the buffer, then use get_all() to retrieve everything as
+    PyTorch tensors for the PPO update step.
+
+    Example:
+        buf = Buffer(step=2048, data={"state": (4,), "action": ()})
+        for _ in range(2048):
+            buf.insert(state=..., action=..., reward=..., ...)
+        tensors = buf.get_all()
+        buf.clear()
+    """
+
+    def __init__(self,
+                 data: dict[str, tuple],
+                 config: TrainConfig):
+        """Initialize pre-allocated arrays.
+
+        Args:
+            step: Maximum number of timesteps (capacity).
+            data: Dict mapping field names to shape tuples (e.g. {"state": (4,), "action": ()}).
+            device: Torch device to allocate tensors on.
+        """
+        self.step = config.rollout_steps
+        self.num_envs = config.num_envs
+        self.slice: int = 0
+        self.data = {
+                name: torch.zeros((self.step, self.num_envs, *shape), dtype = torch.float32, device = config.device)
+                for name, shape in data.items()
+                }
+
+    @property
+    def size(self): return self.slice
+
+    def insert(self, **kwargs):
+        """Insert one timestep of data into the buffer.
+
+        Args:
+            **kwargs: Keyword arguments matching the keys in self.data.
+
+        Raises:
+            ValueError: If the buffer is full.
+            KeyBufferError: If a key doesn't exist in the buffer.
+        """
+        if self.slice >= self.step:
+            raise ValueError(f"Buffer is full (size={self.step}). Cannot insert more data.")
+        for name, val in kwargs.items():
+            if name in self.data:
+                self.data[name][self.slice] = val
+            else:
+                raise KeyBufferError(name, kwargs)
+
+        self.slice += 1
+
+
+    def get_all(self) -> dict[str, torch.Tensor]:
+        """Return all inserted data as a dict of sliced tensors."""
+        return {name: val[:self.slice] for name, val in self.data.items()}
+
+    def clear(self):
+        """Reset the buffer for reuse.
+
+        Resets the slice pointer to 0. Underlying arrays are not zeroed;
+        old data is overwritten by subsequent insert() calls.
+        """
+        self.slice = 0

zerorl/config.py

@@ -0,0 +1,87 @@
+"""Configuration dataclasses for RL training.
+
+Provides AlgoConfig (mutable hyperparameters) and TrainConfig (mutable
+training settings with computed fields).
+"""
+
+import torch
+from dataclasses import dataclass, field
+
+
+@dataclass
+class TrainConfig:
+    """Training configuration with computed fields.
+
+    model_path and num_update are derived in __post_init__(); do not
+    pass them to the constructor.
+
+    Attributes:
+        model_name: Model name (used in the saved file path).
+        model_save_path: Directory for model checkpoints.
+        device: PyTorch device (auto-detects CUDA).
+        model_path: Computed as "{model_save_path}/{model_name}.pt".
+        timestamp: Total environment timesteps for training.
+        rollout_steps: Steps collected before each PPO update.
+        num_update: Computed as timestamp // (rollout_steps * num_envs).
+    """
+    model_name: str
+    project_name: str
+    model_save_path: str = ".checkpoints"
+    timestamp: int = 1_000_000
+    rollout_steps: int = 2048
+    num_envs: int = 1
+    normalize: bool = False
+    profile: bool = False
+
+    device: torch.device = field(init=False)
+    num_update: int = field(init=False)
+    model_path: str = field(init=False)
+
+    def __post_init__(self) -> None:
+        self.device: torch.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model_path = f"{self.model_save_path}/{self.model_name}.pt"
+        self.num_update = self.timestamp // (self.rollout_steps * self.num_envs)
+        if self.num_update <= 0:
+            raise ValueError(f"num_update must be greater than 0, got {self.num_update}")
+
+
+@dataclass(init=False)
+class AlgoConfig:
+    """Mutable algorithm hyperparameters for PPO and off-policy methods.
+
+    Attributes:
+        lr: Learning rate.
+        gamma: Discount factor.
+        batch_size: Minibatch size for PPO updates.
+        gae_lambda: GAE lambda.
+        clip_eps: PPO clipping range.
+        ent_coef: Entropy bonus coefficient.
+        value_coef: Value loss coefficient.
+        epochs: PPO epochs per update.
+        tau: Soft update coefficient (off-policy).
+    """
+
+    lr: float = 3e-4
+    gamma: float = 0.99
+    batch_size: int = 64
+
+    # For on-policy
+    gae_lambda: float = 0.95
+    clip_eps: float = 0.2
+    ent_coef: float = 0.01
+    value_coef: float = 0.5
+    epochs: int = 10
+
+    # For off-policy
+    tau: float = 0.005
+
+    def __init__(self, **kwargs):
+        for key in self.__annotations__:
+            setattr(self, key, getattr(self.__class__, key, None))
+
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def to_dict(self) -> dict:
+        """Return all hyperparameters as a dictionary."""
+        return self.__dict__

zerorl/errors.py

@@ -0,0 +1,59 @@
+"""Custom exceptions for the zerorl project."""
+
+from zerorl.helpers.agent import BaseAgent
+
+
+def assert_agent_contract(agent: BaseAgent, attr_search: dict[str, str]):
+    """
+        Check if attributes exist in Agent
+        agent: Agent
+        attr_search: Key -> Attribute Name
+                     Value -> Message Error
+    """
+    for attr, message in attr_search.items():
+        if not hasattr(agent, attr):
+            raise NotImplementedError(message)
+
+
+class EmptyBufferError(Exception):
+    """Raised when an update is attempted on an insufficiently filled buffer.
+
+    Attributes:
+        current_size: Number of entries currently in the buffer.
+        require_buffer_size: Minimum entries required for an update.
+    """
+
+    def __init__(self, current_size: int, require_buffer_size: int):
+        self.current_size = current_size
+        self.require_buffer_size = require_buffer_size
+
+        self.message = "Training agent flow is incorrect: the buffer is empty"
+        super().__init__(self.message)
+
+    def __str__(self) -> str:
+        """Return a detailed error message with context and suggestion."""
+        suggestion = "Call rollout_phase() before update_weights()"
+        details = f"[Crash Workflow] {self.message}\n"
+        details += f"the current buffer size: {self.current_size}\n"
+        details += f" the minimum size required is: {self.require_buffer_size}\n"
+        details += f"{suggestion}"
+        return details
+
+
+class KeyBufferError(Exception):
+    """Raised when the argument name does not exist in the data buffer.
+
+    Attributes:
+        arg_name: the argument name.
+        data_buffer: the data buffer
+    """
+
+    def __init__(self, arg_name: str, data_buffer: dict[str, object]):
+        self.arg_name = arg_name
+        self.data_buffer = data_buffer
+
+    def __str__(self) -> str:
+        details = f"Key '{self.arg_name}' (returned by agent.get_action) does not exist in the Buffer.\n"
+        details += f"Please ensure your Buffer is initialized with the key '{self.arg_name}'\n"
+        details += f"Current valid keys are: {list(self.data_buffer.keys())}"
+        return details

zerorl/factory.py

@@ -0,0 +1,121 @@
+"""Factory helpers for creating environments, buffers, and agents.
+
+Provides convenience functions that wire together common RL components.
+"""
+
+import numpy as np
+import torch
+from typing import Callable
+from torch import Tensor
+from zerorl.functions import vectorize_env 
+from zerorl.helpers.agent import BaseAgent, eval_action
+from zerorl.helpers.env import BaseEnv
+from zerorl.buffer import Buffer 
+from zerorl.config import TrainConfig 
+from torch import nn
+
+
+def get_env(env_id: str | Callable | BaseEnv, num_envs: int = 1, render_mode: str | None= None):
+    """Create a vectorized environment from a spec, class, or instance.
+
+    Args:
+        env_id: Gymnasium env ID string, BaseEnv class/instance, or callable.
+        num_envs: Number of parallel environments.
+        render_mode: Render mode for the environment.
+
+    Returns:
+        SyncVectorEnv with SAME_STEP autoreset.
+    """
+    return vectorize_env(env_id, num_envs, render_mode)
+
+def get_actor_critic_buffer(state_space: tuple, action_space: tuple, config: TrainConfig):
+    """Create a Buffer with standard PPO field names.
+
+    Args:
+        state_space: Observation shape tuple, e.g. (4,) for a 4-dim vector.
+        action_space: Action shape tuple, e.g. () for discrete or (n,) for continuous.
+        config: Training config providing rollout_steps, num_envs, and device.
+
+    Returns:
+        Buffer pre-allocated with keys: state, action, reward, done, truncated,
+        entropy, value, return, log_prob, advantage.
+    """ 
+    buffer = Buffer(data = {"state": state_space, "action": action_space,
+                            "reward": (), "done": (), "truncated": (), "entropy": (),
+                            "value": (), "return": (), "log_prob": (), "advantage": ()},
+                    config=config)
+    return buffer
+
+class ActorCriticAgent(BaseAgent):
+    """Built-in actor-critic agent with orthogonal initialization.
+
+    Supports both discrete (Categorical) and continuous (Normal) action spaces.
+    Uses a shared 2-layer Tanh MLP feature extractor.
+
+    Args:
+        input_dim: Observation dimension.
+        output_dim: Action dimension (n for discrete, dim for continuous).
+        is_discrete: Whether the action space is discrete.
+        hidden_dim: Hidden layer size (default 64).
+    """
+    def __init__(self, input_dim: int, output_dim: int, is_discrete: bool, hidden_dim: int = 64):
+        super().__init__()
+
+        self.is_discrete = is_discrete
+        self.hidden_dim = hidden_dim
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        
+        # Feature Extractor
+        self.extract_layer = nn.Sequential(
+                nn.Linear(self.input_dim, self.hidden_dim),
+                nn.Tanh(),
+                nn.Linear(self.hidden_dim, self.hidden_dim),
+                nn.Tanh()
+                )
+        # Actor
+        self.actor = nn.Linear(self.hidden_dim, self.output_dim)
+        # Critic
+        self.critic = nn.Linear(self.hidden_dim,  1)
+
+        if not is_discrete:
+            self.log_std = nn.Parameter(torch.zeros(output_dim))
+
+        self.apply(self._orthogonal_init)
+
+
+    def _orthogonal_init(self, module: nn.Module):
+        """Apply orthogonal weight initialization with gain based on layer role."""
+        if isinstance(module, nn.Linear):
+            if module.out_features == self.hidden_dim:
+                nn.init.orthogonal_(module.weight, gain=np.sqrt(2))
+            elif module.out_features == 1:
+                nn.init.orthogonal_(module.weight, gain=1.0)
+            else:
+                nn.init.orthogonal_(module.weight, gain=0.01)
+                
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0.0)
+
+    def forward(self, state: Tensor):
+        """Forward pass returning (logits, value)."""
+        x = self.extract_layer(state)
+        logits = self.actor(x)
+        value = self.critic(x)
+        return (logits, value)
+
+    def build_distribution(self, logits: torch.Tensor):
+        """Build a torch distribution from logits (Categorical or Normal)."""
+        if self.is_discrete:
+            return torch.distributions.Categorical(logits=logits)
+        log_std_clamped = torch.clamp(self.log_std, min=-3.0, max=1.0)
+        std = log_std_clamped.exp().expand_as(logits)
+        return torch.distributions.Normal(logits, std)
+    
+    def get_action(self, state: torch.Tensor, action: torch.Tensor | None = None):
+        """Sample or evaluate an action, returning action, log_prob, entropy, value."""
+        logits, value = self.forward(state)
+        dist = self.build_distribution(logits)
+        if action is None: action = dist.sample()
+        log_prob, dist_entropy = eval_action(dist, action)
+        return {"action": action, "log_prob": log_prob, "entropy":dist_entropy, "value":value}