def td_lambda(
v_tm1: ArrayLike,
r_t: ArrayLike,
discount_t: ArrayLike,
v_t: ArrayLike,
lambda_: ArrayOrScalar,
) -> ArrayLike:
"""Calculates the TD(lambda) temporal difference error.
See "Reinforcement Learning: An Introduction" by Sutton and Barto.
(http://incompleteideas.net/book/ebook/node74.html).
Args:
v_tm1: sequence of state values at time t-1.
r_t: sequence of rewards at time t.
discount_t: sequence of discounts at time t.
v_t: sequence of state values at time t.
lambda_: mixing parameter lambda, either a scalar or a sequence.
Returns:
TD(lambda) temporal difference error.
"""
base.rank_assert([v_tm1, r_t, discount_t, v_t, lambda_],
[1, 1, 1, 1, [0, 1]])
base.type_assert([v_tm1, r_t, discount_t, v_t, lambda_], float)
target_tm1 = multistep.lambda_returns(r_t, discount_t, v_t, lambda_)
return jax.lax.stop_gradient(target_tm1) - v_tm1
def td_lambda(
v_tm1: Array,
r_t: Array,
discount_t: Array,
v_t: Array,
lambda_: Numeric,
stop_target_gradients: bool = True,
) -> Array:
"""Calculates the TD(lambda) temporal difference error.
See "Reinforcement Learning: An Introduction" by Sutton and Barto.
(http://incompleteideas.net/book/ebook/node74.html).
Args:
v_tm1: sequence of state values at time t-1.
r_t: sequence of rewards at time t.
discount_t: sequence of discounts at time t.
v_t: sequence of state values at time t.
lambda_: mixing parameter lambda, either a scalar or a sequence.
stop_target_gradients: bool indicating whether or not to apply stop gradient
to targets.
Returns:
TD(lambda) temporal difference error.
"""
chex.assert_rank([v_tm1, r_t, discount_t, v_t, lambda_],
[1, 1, 1, 1, {0, 1}])
chex.assert_type([v_tm1, r_t, discount_t, v_t, lambda_], float)
target_tm1 = multistep.lambda_returns(r_t, discount_t, v_t, lambda_)
target_tm1 = jax.lax.select(stop_target_gradients,
jax.lax.stop_gradient(target_tm1), target_tm1)
return target_tm1 - v_tm1
def test_reduces_to_lambda_returns(self):
"""Test function is the same as lambda_returns when n is sequence length."""
lambda_t = 0.75
n = len(self.r_t[0])
expected = multistep.lambda_returns(self.r_t[0], self.discount_t[0],
self.v_t[0], lambda_t)
actual = multistep.n_step_bootstrapped_returns(self.r_t[0],
self.discount_t[0],
self.v_t[0], n,
lambda_t)
np.testing.assert_allclose(expected, actual, rtol=1e-5)
def sarsa_lambda(
q_tm1: ArrayLike,
a_tm1: ArrayLike,
r_t: ArrayLike,
discount_t: ArrayLike,
q_t: ArrayLike,
a_t: ArrayLike,
lambda_: ArrayOrScalar,
stop_target_gradients: bool = True,
) -> ArrayLike:
"""Calculates the SARSA(lambda) temporal difference error.
See "Reinforcement Learning: An Introduction" by Sutton and Barto.
(http://incompleteideas.net/book/ebook/node77.html).
Args:
q_tm1: sequence of Q-values at time t-1.
a_tm1: sequence of action indices at time t-1.
r_t: sequence of rewards at time t.
discount_t: sequence of discounts at time t.
q_t: sequence of Q-values at time t.
a_t: sequence of action indices at time t.
lambda_: mixing parameter lambda, either a scalar or a sequence.
stop_target_gradients: bool indicating whether or not to apply stop gradient
to targets.
Returns:
SARSA(lambda) temporal difference error.
"""
base.rank_assert([q_tm1, a_tm1, r_t, discount_t, q_t, a_t, lambda_],
[2, 1, 1, 1, 2, 1, [0, 1]])
base.type_assert([q_tm1, a_tm1, r_t, discount_t, q_t, a_t, lambda_],
[float, int, float, float, float, int, float])
qa_tm1 = base.batched_index(q_tm1, a_tm1)
qa_t = base.batched_index(q_t, a_t)
target_tm1 = multistep.lambda_returns(r_t, discount_t, qa_t, lambda_)
if stop_target_gradients:
target_tm1 = jax.lax.stop_gradient(target_tm1)
return target_tm1 - qa_tm1
def q_lambda(
q_tm1: Array,
a_tm1: Array,
r_t: Array,
discount_t: Array,
q_t: Array,
lambda_: Numeric,
stop_target_gradients: bool = True,
) -> Array:
"""Calculates Peng's or Watkins' Q(lambda) temporal difference error.
See "Reinforcement Learning: An Introduction" by Sutton and Barto.
(http://incompleteideas.net/book/ebook/node78.html).
Args:
q_tm1: sequence of Q-values at time t-1.
a_tm1: sequence of action indices at time t-1.
r_t: sequence of rewards at time t.
discount_t: sequence of discounts at time t.
q_t: sequence of Q-values at time t.
lambda_: mixing parameter lambda, either a scalar (e.g. Peng's Q(lambda)) or
a sequence (e.g. Watkin's Q(lambda)).
stop_target_gradients: bool indicating whether or not to apply stop gradient
to targets.
Returns:
Q(lambda) temporal difference error.
"""
chex.assert_rank([q_tm1, a_tm1, r_t, discount_t, q_t, lambda_],
[2, 1, 1, 1, 2, {0, 1}])
chex.assert_type([q_tm1, a_tm1, r_t, discount_t, q_t, lambda_],
[float, int, float, float, float, float])
qa_tm1 = base.batched_index(q_tm1, a_tm1)
v_t = jnp.max(q_t, axis=-1)
target_tm1 = multistep.lambda_returns(r_t, discount_t, v_t, lambda_)
target_tm1 = jax.lax.select(stop_target_gradients,
jax.lax.stop_gradient(target_tm1), target_tm1)
return target_tm1 - qa_tm1