When using TD(λ), how do you calculate the eligibility trace per input & weight of a neural network neuron?

Question

I have a Neural Network, each Neuron is made up of inputs, weights, and output. I have potentially multiple hidden layers. The activation function executed against the output is not known by the Neuron.

I would like to use TD(λ) to back-propagate errors through the network as it explores options. My understanding is that this is forward looking TD(λ) because I won't know the error until I reach a terminal state, and so an eligibility trace needs to be kept for each input+weight combination as I back-propagate the error between the NNs new output based on the state-change from the last prediction and the output from the last prediction.

To try and modularise my code as much as possible, the neuron won't know the loss function, but will instead be given the error as a derivative of whatever the activation function of its output was. It also won't know if it's in a hidden layer or not, it will just have inputs, weights, and an output

For example:

So when each Neuron is back-propagating the error signal from all its output connections (summed before it receives it), how do I calculate the eligibility trace?

Answer 1

Eligibility traces are not very common in deep-RL, but I suppose that if your network is not huge they might work. I think you are describing the backward view of eligibility traces since you are saying you want to keep a "trace" of each state visited before reaching the terminal state.

In this case, computing and updating the traces to later combine them with the TD-Error could be done as follows. You would need to initialize a $\textbf{Z}$ matrix with the same dimensions as the weights $(\textbf{W})$ for each layer of your network. Then, at each time step you update the trace by simply:

$$ \textbf{Z}_{ij} \leftarrow \gamma\lambda \textbf{Z}{ij} + \nabla_{\textbf{W}_{ij}}O(i) $$

where $O(i)$ is the output of the network. Assuming that $\textbf{O}$ is the last layer and either O1 is selected or O2 (but not both, as in a policy network) this is the rule. $\gamma$ is the usual discount factor and $\lambda$ the trace decay parameter. Once you have this then update the weights with the usual TD learning rule:

$$ \textbf{W}_{ij} \leftarrow \textbf{W}{ij} + \delta\textbf{Z}_{ij}$$

The TD-Error is whatever you prefer (i.e. Actor-Critic or SARSA). To understand better this combined rule, set $\lambda$ to be zero and see how you recover the usual TD(0) learning rule. I imagine you are familiar with Sutton&Barto 2018 fantastic book in RL (not deep RL). This topic is explained in chapter 12.

Recall that at each state (or time step) all the information to compute these rules is available thanks to the forward pass from the input to the output.

Hope this helped.