lit_sb3.agents

source package lit_sb3.agents

Classes

• BaseAgent — Abstract base for RL agents.

• OffPolicyAgent — Abstract base for off-policy agents.

• OnPolicyAgent — Abstract base for on-policy 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.

source class OffPolicyAgent(policy: BasePolicy)

Bases : BaseAgent

Abstract base for off-policy agents.

Subclasses must implement both :meth:configure_optimizers and :meth:compute_loss.

Methods

• configure_optimizers — Return the list of optimizers (varies: SAC=3, TD3=2, DQN=1).

• compute_loss — Compute off-policy loss for one replay-buffer batch.

source method OffPolicyAgent.configure_optimizers() → OptimizerLRScheduler

Return the list of optimizers (varies: SAC=3, TD3=2, DQN=1).

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

Compute off-policy loss for one replay-buffer batch.

source class OnPolicyAgent(policy: BasePolicy)

Bases : BaseAgent

Abstract base for on-policy agents.

Provides a concrete :meth:configure_optimizers (single policy optimizer). Subclasses must implement :meth:compute_loss.

Methods

• configure_optimizers — Return the single policy optimizer.

• compute_loss — Compute on-policy loss for one rollout mini-batch.

source method OnPolicyAgent.configure_optimizers() → OptimizerLRScheduler

Return the single policy optimizer.

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

Compute on-policy loss for one rollout mini-batch.