lit_sb3.agents.base

source module lit_sb3.agents.base

BaseAgent — abstract base class for RL agents.

An *agent* owns the policy (neural network) and the learning objective

• compute_loss(batch) — pure loss computation, returns a dict of tensors.
• configure_optimizers() — returns the list of optimizers for the policy.
• get_action(obs) — action selection (overridable for ε-greedy, noise, …).

The algorithm (BaseAlgorithm) remains the LightningModule orchestrator that owns the environment, replay/rollout buffer, data collection loop, and the Lightning hooks (training_step, log_dict, manual_backward). Agents never call Lightning methods directly.

Classes

• BaseAgent — Abstract base for RL agents.

source class BaseAgent(policy: BasePolicy)

Bases : ABC

Abstract base for RL agents.

Subclasses must implement :meth:configure_optimizers and :meth:compute_loss. :meth:get_action and :meth:set_training_mode have sensible defaults that delegate to self.policy.

policy must be set by the subclass __init__ (or _setup) before any method is called.

Methods

• configure_optimizers — Return the list of optimizers used by this agent.

• compute_loss — Compute and return losses for batch.

• get_action — Select an action given observation.

• set_training_mode — Toggle train/eval mode on the underlying policy.

source method BaseAgent.configure_optimizers() → OptimizerLRScheduler

Return the list of optimizers used by this agent.

Called once by the algorithm's configure_optimizers Lightning hook.

source method BaseAgent.compute_loss(batch: Any) → dict[str, torch.Tensor]

Compute and return losses for batch.

Returns a dict mapping loss names to scalar tensors. The dict must contain at least a "loss" key (the total loss to differentiate). All other entries are treated as diagnostic scalars for logging.

The algorithm is responsible for calling manual_backward and log_dict — this method must not do so itself.

source method BaseAgent.get_action(observation: np.ndarray | dict[str, np.ndarray], state: tuple[np.ndarray, ...] | None = None, episode_start: np.ndarray | None = None, deterministic: bool = False) → tuple[np.ndarray, tuple[np.ndarray, ...] | None]

Select an action given observation.

Default implementation delegates to self.policy.predict. Override for exploration strategies (ε-greedy, action noise, …).

source method BaseAgent.set_training_mode(mode: bool) → None

Toggle train/eval mode on the underlying policy.