RL sampler (experimental)

ResidualBanditSampler

ResidualBanditSampler(
    domain,
    time_domain,
    n_x_cells=16,
    n_t_cells=8,
    policy_fraction=0.5,
    temperature=0.5,
    probe_cells=32,
    probe_points=8,
    probe_every=1,
    agent=None,
    device=None,
)

Collocation sampler whose cell preferences are learned by a DQN bandit.

Parameters:

Name	Type	Description	Default
domain	tuple[float, float]	(x_lo, x_hi) spatial bounds (1-D PDEs only).	required
time_domain	tuple[float, float]	(t_lo, t_hi) temporal bounds.	required
n_x_cells	int	Number of cells along x.	16
n_t_cells	int	Number of cells along t.	8
policy_fraction	float	Share of each batch drawn from the policy; the rest is uniform over the full domain.	0.5
temperature	float	Softmax temperature on the Q-values.	0.5
probe_cells	int	Cells probed per call to collect rewards.	32
probe_points	int	Residual evaluations per probed cell.	8
probe_every	int	Probe (and train the agent) on every probe_every-th call to :meth:sample; other calls reuse the last state. Raise it when sampling every optimizer step, as the trainer does.	1
agent	RLAgent | None	Optional pre-built agent; must have state_dim=4 and action_dim = n_x_cells * n_t_cells.	None
device	device | None	Torch device.	None

Source code in pinnrl/rl/bandit_sampler.py

def __init__(
    self,
    domain: tuple[float, float],
    time_domain: tuple[float, float],
    n_x_cells: int = 16,
    n_t_cells: int = 8,
    policy_fraction: float = 0.5,
    temperature: float = 0.5,
    probe_cells: int = 32,
    probe_points: int = 8,
    probe_every: int = 1,
    agent: RLAgent | None = None,
    device: torch.device | None = None,
):
    self.device = device or torch.device("cpu")
    self.x_lo, self.x_hi = float(domain[0]), float(domain[1])
    self.t_lo, self.t_hi = float(time_domain[0]), float(time_domain[1])
    self.n_x_cells = n_x_cells
    self.n_t_cells = n_t_cells
    self.n_cells = n_x_cells * n_t_cells
    self.policy_fraction = policy_fraction
    self.temperature = temperature
    self.probe_cells = probe_cells
    self.probe_points = probe_points
    self.probe_every = max(1, probe_every)
    self._calls = 0
    self._state: torch.Tensor | None = None
    self.agent = agent or RLAgent(
        state_dim=self.STATE_DIM,
        action_dim=self.n_cells,
        hidden_dim=64,
        learning_rate=2e-3,
        gamma=0.0,
        epsilon_start=1.0,
        epsilon_end=0.05,
        epsilon_decay=0.998,
        memory_size=4000,
        batch_size=64,
        target_update=100,
        device=self.device,
    )
    if self.agent.action_dim != self.n_cells:
        raise ValueError(
            f"agent.action_dim={self.agent.action_dim} does not match "
            f"n_x_cells * n_t_cells = {self.n_cells}"
        )

probe

probe(pde, model, progress)

Probe random cells, push rewards to the agent, and return the state.

probe_points random points in each of min(probe_cells, n_cells) distinct cells are scored by |residual|. The state is [log10 mean, log10 max, log10 std, progress] over all probed values (the mean is the mean of the cell means); each cell's reward is its mean divided by that overall mean. Each transition is stored with next_state = state, which gamma = 0 ignores.

Source code in pinnrl/rl/bandit_sampler.py

def probe(self, pde, model: torch.nn.Module, progress: float) -> torch.Tensor:
    """Probe random cells, push rewards to the agent, and return the state.

    ``probe_points`` random points in each of ``min(probe_cells, n_cells)``
    distinct cells are scored by ``|residual|``. The state is
    ``[log10 mean, log10 max, log10 std, progress]`` over all probed
    values (the mean is the mean of the cell means); each cell's reward
    is its mean divided by that overall mean. Each transition is stored
    with ``next_state = state``, which ``gamma = 0`` ignores.
    """
    cells = torch.randperm(self.n_cells, device=self.device)[: self.probe_cells]
    x, t = self._points_in_cells(cells.repeat_interleave(self.probe_points))
    residual = pde.compute_residual(model, x.requires_grad_(True), t.requires_grad_(True))
    res_abs = residual.detach().abs().reshape(len(cells), self.probe_points)

    cell_means = res_abs.mean(dim=1)
    overall = cell_means.mean().clamp_min(1e-12)
    state = torch.tensor(
        [
            torch.log10(overall).item(),
            torch.log10(res_abs.max().clamp_min(1e-12)).item(),
            torch.log10(res_abs.std().clamp_min(1e-12)).item(),
            float(progress),
        ],
        device=self.device,
    )
    for cell, reward in zip(cells.tolist(), (cell_means / overall).tolist(), strict=False):
        self.agent.update(state, cell, reward, state, False)
    return state

distribution

distribution(state)

ε-mixture of softmax(Q / temperature) and the uniform distribution.

The Q-network is evaluated in eval mode (dropout off), so the same state always gives the same distribution.

Source code in pinnrl/rl/bandit_sampler.py

def distribution(self, state: torch.Tensor) -> torch.Tensor:
    """ε-mixture of ``softmax(Q / temperature)`` and the uniform distribution.

    The Q-network is evaluated in eval mode (dropout off), so the same
    state always gives the same distribution.
    """
    net = self.agent.policy_net
    was_training = net.training
    net.eval()
    with torch.no_grad():
        q = net(state.unsqueeze(0)).squeeze(0)
    net.train(was_training)
    eps = float(self.agent.epsilon)
    return (1.0 - eps) * torch.softmax(q / self.temperature, dim=0) + eps / self.n_cells

sample

sample(pde, model, num_points, progress=0.0)

Return num_points collocation points, as (x, t) columns.

Probes and updates the agent first on every probe_every-th call (and on the first); other calls reuse the last state. progress in [0, 1] is the training-progress entry of the state.

Source code in pinnrl/rl/bandit_sampler.py

def sample(
    self, pde, model: torch.nn.Module, num_points: int, progress: float = 0.0
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return ``num_points`` collocation points, as ``(x, t)`` columns.

    Probes and updates the agent first on every ``probe_every``-th call
    (and on the first); other calls reuse the last state. ``progress`` in
    [0, 1] is the training-progress entry of the state.
    """
    if self._state is None or self._calls % self.probe_every == 0:
        self._state = self.probe(pde, model, progress)
    self._calls += 1
    state = self._state
    n_policy = int(round(num_points * self.policy_fraction))
    if n_policy > 0:
        cells = torch.multinomial(self.distribution(state), n_policy, replacement=True)
    else:
        cells = torch.empty(0, dtype=torch.long, device=self.device)
    x_p, t_p = self._points_in_cells(cells)
    x_u, t_u = self._uniform(num_points - n_policy)
    return torch.cat([x_p, x_u]).detach(), torch.cat([t_p, t_u]).detach()

RLAgent

RLAgent(
    state_dim,
    action_dim,
    hidden_dim,
    learning_rate=0.0001,
    gamma=0.99,
    epsilon_start=1.0,
    epsilon_end=0.01,
    epsilon_decay=0.995,
    memory_size=10000,
    batch_size=64,
    target_update=100,
    device=None,
)

DQN agent with experience replay and an ε schedule that decays per update.

Initialize RL agent.

Parameters:

Name	Type	Description	Default
state_dim	int	State dimension.	required
action_dim	int	Number of discrete actions.	required
hidden_dim	int	Hidden layer dimension.	required
learning_rate	float	Learning rate.	0.0001
gamma	float	Discount factor (0 turns the DQN into a contextual bandit).	0.99
epsilon_start	float	Initial exploration rate.	1.0
epsilon_end	float	Final exploration rate.	0.01
epsilon_decay	float	Multiplicative ε decay applied on every :meth:update.	0.995
memory_size	int	Size of replay buffer.	10000
batch_size	int	Batch size for training.	64
target_update	int	Updates between target-network syncs.	100
device	device | None	Device to place the networks on.	None

Source code in pinnrl/rl/rl_agent.py

def __init__(
    self,
    state_dim: int,
    action_dim: int,
    hidden_dim: int,
    learning_rate: float = 0.0001,
    gamma: float = 0.99,
    epsilon_start: float = 1.0,
    epsilon_end: float = 0.01,
    epsilon_decay: float = 0.995,
    memory_size: int = 10000,
    batch_size: int = 64,
    target_update: int = 100,
    device: torch.device | None = None,
):
    """Initialize RL agent.

    Args:
        state_dim: State dimension.
        action_dim: Number of discrete actions.
        hidden_dim: Hidden layer dimension.
        learning_rate: Learning rate.
        gamma: Discount factor (``0`` turns the DQN into a contextual bandit).
        epsilon_start: Initial exploration rate.
        epsilon_end: Final exploration rate.
        epsilon_decay: Multiplicative ε decay applied on every :meth:`update`.
        memory_size: Size of replay buffer.
        batch_size: Batch size for training.
        target_update: Updates between target-network syncs.
        device: Device to place the networks on.
    """
    self.device = device or torch.device("cpu")
    self.state_dim = state_dim
    self.action_dim = action_dim
    self.hidden_dim = hidden_dim
    self.learning_rate = learning_rate
    self.gamma = gamma
    self.epsilon = epsilon_start
    self.epsilon_end = epsilon_end
    self.epsilon_decay = epsilon_decay
    self.batch_size = batch_size
    self.target_update = target_update

    self.policy_net = DQNNetwork(state_dim, action_dim, hidden_dim).to(self.device)
    self.target_net = DQNNetwork(state_dim, action_dim, hidden_dim).to(self.device)
    self.target_net.load_state_dict(self.policy_net.state_dict())
    self.optimizer = optim.Adam(self.policy_net.parameters(), lr=learning_rate)
    self.memory = ReplayBuffer(memory_size)
    self.steps = 0

update

update(state, action, reward, next_state, done)

Store a transition, train on a replayed batch, and decay ε.

Training starts once the buffer holds batch_size transitions; the target network is synced every target_update calls.

Parameters:

Name	Type	Description	Default
state	Tensor	Current state.	required
action	int	Action taken.	required
reward	float	Reward received.	required
next_state	Tensor	Next state.	required
done	bool	Whether the episode ended.	required

Source code in pinnrl/rl/rl_agent.py

def update(
    self,
    state: torch.Tensor,
    action: int,
    reward: float,
    next_state: torch.Tensor,
    done: bool,
):
    """Store a transition, train on a replayed batch, and decay ε.

    Training starts once the buffer holds ``batch_size`` transitions; the
    target network is synced every ``target_update`` calls.

    Args:
        state: Current state.
        action: Action taken.
        reward: Reward received.
        next_state: Next state.
        done: Whether the episode ended.
    """
    self.memory.push(state, action, reward, next_state, done)
    self.steps += 1
    if self.steps % self.target_update == 0:
        self.target_net.load_state_dict(self.policy_net.state_dict())
    if len(self.memory) >= self.batch_size:
        self._train()
    self.epsilon = max(self.epsilon_end, self.epsilon * self.epsilon_decay)

DQNNetwork

DQNNetwork(
    state_dim,
    action_dim,
    hidden_dim,
    num_layers=3,
    dropout=0.1,
)

Bases: Module

MLP mapping a state to one Q-value per action (Linear-LayerNorm-ReLU-Dropout blocks).

Initialize DQN network.

Parameters:

Name	Type	Description	Default
state_dim	int	State dimension.	required
action_dim	int	Action dimension.	required
hidden_dim	int	Hidden layer dimension.	required
num_layers	int	Number of linear layers, input and output included.	3
dropout	float	Dropout rate after each hidden block.	0.1

Source code in pinnrl/rl/rl_agent.py

def __init__(
    self,
    state_dim: int,
    action_dim: int,
    hidden_dim: int,
    num_layers: int = 3,
    dropout: float = 0.1,
):
    """Initialize DQN network.

    Args:
        state_dim: State dimension.
        action_dim: Action dimension.
        hidden_dim: Hidden layer dimension.
        num_layers: Number of linear layers, input and output included.
        dropout: Dropout rate after each hidden block.
    """
    super().__init__()

    # Build network
    layers = []

    # Input layer
    layers.append(
        nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
        )
    )

    # Hidden layers
    for _ in range(num_layers - 2):
        layers.append(
            nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim),
                nn.LayerNorm(hidden_dim),
                nn.ReLU(),
                nn.Dropout(dropout),
            )
        )

    # Output layer
    layers.append(nn.Linear(hidden_dim, action_dim))

    self.layers = nn.Sequential(*layers)

    # Initialize weights
    self._init_weights()

forward

forward(x)

Forward pass through the network.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor.	required

Returns:

Type	Description
Tensor	Q-values.

Source code in pinnrl/rl/rl_agent.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass through the network.

    Args:
        x: Input tensor.

    Returns:
        Q-values.
    """
    return self.layers(x)

ReplayBuffer

ReplayBuffer(capacity)

Experience replay buffer for training stability.

Initialize replay buffer.

Parameters:

Name	Type	Description	Default
capacity	int	Maximum number of experiences to store.	required

Source code in pinnrl/rl/rl_agent.py

def __init__(self, capacity: int):
    """Initialize replay buffer.

    Args:
        capacity: Maximum number of experiences to store.
    """
    self.buffer = deque(maxlen=capacity)

push

push(state, action, reward, next_state, done)

Add experience to buffer.

Parameters:

Name	Type	Description	Default
state	Tensor	Current state.	required
action	int	Action taken.	required
reward	float	Reward received.	required
next_state	Tensor	Next state.	required
done	bool	Whether episode is done.	required

Source code in pinnrl/rl/rl_agent.py

def push(
    self,
    state: torch.Tensor,
    action: int,
    reward: float,
    next_state: torch.Tensor,
    done: bool,
):
    """Add experience to buffer.

    Args:
        state: Current state.
        action: Action taken.
        reward: Reward received.
        next_state: Next state.
        done: Whether episode is done.
    """
    self.buffer.append((state, action, reward, next_state, done))

sample

sample(batch_size)

Sample a batch of experiences without replacement.

Indices come from torch's generator, so torch.manual_seed makes the draw (and a bandit run) reproducible.

Parameters:

Name	Type	Description	Default
batch_size	int	Size of batch to sample.	required

Returns:

Type	Description
list[tuple]	List of experiences.

Raises:

Type	Description
ValueError	If batch_size exceeds the stored experiences.

Source code in pinnrl/rl/rl_agent.py

def sample(self, batch_size: int) -> list[tuple]:
    """Sample a batch of experiences without replacement.

    Indices come from torch's generator, so ``torch.manual_seed`` makes
    the draw (and a bandit run) reproducible.

    Args:
        batch_size: Size of batch to sample.

    Returns:
        List of experiences.

    Raises:
        ValueError: If ``batch_size`` exceeds the stored experiences.
    """
    if batch_size > len(self.buffer):
        raise ValueError(
            f"batch_size {batch_size} exceeds the {len(self.buffer)} stored experiences"
        )
    return [self.buffer[i] for i in torch.randperm(len(self.buffer))[:batch_size].tolist()]

__len__

__len__()

Get current size of buffer.

Source code in pinnrl/rl/rl_agent.py

def __len__(self) -> int:
    """Get current size of buffer."""
    return len(self.buffer)