REINFORCE

REINFORCE

The implementation follows ordinary episodic REINFORCE. This page only describes the details that matter when a spiking policy preserves neuron state between environment steps. The agent is implemented in src/spiking_rl_lab/agents/reinforce.py.

Adapting REINFORCE to spiking policies

A spiking policy depends on both the current observation and the neuron state:

$$ s_{t+1}=F_\theta(x_t,s_t), \qquad \pi_t=\Pi(H_\theta(x_t,s_t)). $$

Here $x_t$ is the processed observation, $s_t$ is the neuron state, and $\theta$ denotes trainable policy parameters. $F_\theta$ updates the neuron state, $H_\theta$ produces the network output, and $\Pi$ converts that output into an action distribution $\pi_t$.

Each environment step advances the network by exactly one LIF or LI update. The agent does not create an inner spike-time window and does not perform rate or population decoding. LIF nodes expose spikes, while LI or linear nodes can provide a continuous readout after a spiking hidden layer.

Interaction with the environment runs under torch.no_grad(). At the start of a rollout, the agent saves the network state. During optimization it restores that state and sequentially replays the stored observations and actions with gradients enabled. The usual REINFORCE loss is then evaluated on the reconstructed temporal graph:

$$ \mathcal L=-\frac{1}{TN}\sum_{t,n}G_t^{(n)}\log\pi_\theta\left(a_t^{(n)}\mid x_t^{(n)},s_t^{(n)}\right). $$

Here $T$ is the rollout length, $N$ is the number of parallel environments, $G_t^{(n)}$ is the discounted return, and $a_t^{(n)}$ is the action sampled during collection. The log-probability is recomputed during replay from the stored observation, action, and reconstructed neuron state.

Replay is important for spiking networks because the action probability at step $t$ depends on membrane and synaptic states accumulated at earlier steps. Surrogate gradients configured by the LIF node's method and alpha parameters make backpropagation through spike generation possible.

State boundaries

State is stored separately for every parallel environment. When an environment terminates or is truncated, only its batch row is reset to the node's resting state:

$$ s_{t+1}^{(n)} \leftarrow \begin{cases} s_{\mathrm{rest}}, & d_t^{(n)}=1,\\ s_{t+1}^{(n)}, & d_t^{(n)}=0. \end{cases} $$

Here $d_t^{(n)}$ indicates termination or truncation of environment $n$, while $s_{\mathrm{rest}}$ is the initial resting state of the corresponding node.

The same reset is applied during collection and replay, preventing membrane voltage and synaptic current from crossing episode boundaries. Clearing rollout memory does not reset the live state, so an unfinished episode can continue across rollouts.

During replay, node states are detached every sequence_length transitions:

$$ s_{kL}\leftarrow \operatorname{stopgrad}(s_{kL}), \qquad L=\texttt{sequence_length}. $$

Here $L$ is the TBPTT window length and $kL$ denotes each window boundary. $\operatorname{stopgrad}$ keeps the current numerical state but cuts its connection to the preceding computation graph.

Increasing sequence_length allows longer temporal credit assignment at a higher memory cost.

Practical constraints

• The final network width must match flatdim(action_space).
• A final LIF node uses binary spikes directly as categorical logits or Gaussian means. An LI or linear readout is usually less restrictive.
• The categorical and Gaussian policies provide trainable log-probabilities. The deterministic policy produces zero log-probability and therefore cannot train through REINFORCE alone.
• Returns stop at episode boundaries and at the end of the collected rollout. There is no value-function bootstrap.

The main SNN-specific parameter is agent.params.sequence_length. Node dynamics and surrogate gradients are configured in the lif and li node definitions. See Configuration and experiments for composition and override syntax.