lit_sb3.agents.off_policy
[docs]
module
lit_sb3.agents.off_policy
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32 | """OffPolicyAgent — base for replay-buffer-based agents (SAC, TD3, DQN, …).
Off-policy agents learn from transitions stored in a replay buffer.
``configure_optimizers`` is kept abstract because the optimizer set varies
significantly across algorithms (SAC has 3, TD3 has 2, DQN has 1).
"""
from __future__ import annotations
from abc import abstractmethod
from typing import TYPE_CHECKING, Any
from lightning.pytorch.utilities.types import OptimizerLRScheduler
import torch
from lit_sb3.agents.base import BaseAgent
class OffPolicyAgent(BaseAgent):
"""Abstract base for off-policy agents.
Subclasses must implement both :meth:`configure_optimizers` and
:meth:`compute_loss`.
"""
@abstractmethod
def configure_optimizers(self) -> OptimizerLRScheduler:
"""Return the list of optimizers (varies: SAC=3, TD3=2, DQN=1)."""
@abstractmethod
def compute_loss(self, batch: Any) -> dict[str, torch.Tensor]:
"""Compute off-policy loss for one replay-buffer batch."""
|