SARSATrace#

class SARSATrace(learning_rate=None, eligibility_trace=None, model_free_values=None, second_stage_values=None, starting_state=None, action=None, reached_second_stage=None, reward=None, **kwargs)[source]#

Bases: object

SARSA learning rule with an eligibility trace for a two-stage Markov decision task (Sutton & Barto, 2018; Kool et al., 2016).

Parameters:

learning_ratefloat: The learning rate, \(\alpha\).
eligibility_tracefloat: The eligibility trace decay, \(\lambda\). It determines how much of the second-stage prediction error is carried back to the first-stage model-free value.
model_free_valuesndarray: The first-stage model-free Q-values, a 2D array of shape (n_states, n_actions). The array is not modified in-place.
second_stage_valuesndarray: The second-stage state values, a 1D array of shape (n_second_stage_states,). The array is not modified in-place.
starting_stateint: The first-stage starting state (0-indexed).
actionint: The chosen first-stage action (0-indexed).
reached_second_stageint: The second-stage state reached after the transition (0-indexed).
rewardfloat: The reward received in the second-stage state.

Attributes:

stage1_prediction_errorfloat: The first-stage prediction error, \(\delta_1\), set after calling compute().
stage2_prediction_errorfloat: The second-stage prediction error, \(\delta_2\), set after calling compute().
model_free_deltandarray: The change to add to model_free_values, set after calling compute().
planet_value_deltandarray: The change to add to second_stage_values, set after calling compute().

Notes

The rule computes two coupled prediction errors. The first-stage prediction error compares the value of the reached second-stage state with the model-free value of the chosen first-stage action,

\[\delta_1 = Q_2(s_2) - Q_{MF}(s_1, a),\]

and the second-stage prediction error compares the reward with the value of the reached second-stage state,

\[\delta_2 = r - Q_2(s_2).\]

The values are then updated as

\[\Delta Q_{MF}(s_1, a) = \alpha \delta_1 + \lambda \alpha \delta_2,\]

\[\Delta Q_2(s_2) = \alpha \delta_2,\]

where the eligibility trace \(\lambda\) carries the second-stage prediction error back to the first-stage choice. All other values remain unchanged.

References

Kool, W., Cushman, F. A., & Gershman, S. J. (2016). When does model-based control pay off? PLoS Computational Biology, 12(8), e1005090.

Sutton, R. S., & Barto, A. G. (2018). Reinforcement learning: An introduction (Second edition). The MIT Press.

Examples

>>> import numpy as np
>>> from cpm.models.learning import SARSATrace
>>> model_free_values = np.zeros((2, 2))
>>> second_stage_values = np.array([4.5, 4.5])
>>> sarsa = SARSATrace(
...     learning_rate=0.5, eligibility_trace=0.6,
...     model_free_values=model_free_values, second_stage_values=second_stage_values,
...     starting_state=0, action=1, reached_second_stage=0, reward=7.0,
... )
>>> model_free_delta, planet_value_delta = sarsa.compute()
>>> model_free_delta
array([[0., 3.],
       [0., 0.]])
>>> planet_value_delta
array([1.25, 0.  ])

compute()[source]#

Compute the changes in the model-free and second-stage values.

Returns:

model_free_deltanumpy.ndarray: The change to add to the model-free values. Non-zero only at [starting_state, action].
planet_value_deltanumpy.ndarray: The change to add to the second-stage values. Non-zero only at [reached_second_stage].