class_docstr: An online Q-lambda agent.
class_docstr_tok: ["An","online","Q-lambda","agent","."]
docstr: Calculates Q-lambda loss given parameters, actions and timesteps.
docstr_tok: ["Calculates","Q-lambda","loss","given","parameters",",","actions","and","timesteps","."]
code: def _loss(self, params, actions, timesteps):
      """Calculates Q-lambda loss given parameters, actions and timesteps."""
      network_apply_sequence = jax.vmap(
          self._network.apply, in_axes=(None, 0)
      )
      q = network_apply_sequence(params, timesteps.observation)
  
      # Use a mask since the sequence could cross episode boundaries.
      mask = jnp.not_equal(
          timesteps.step_type, int(dm_env.StepType.LAST)
      )
      a_tm1 = actions[1:]
      r_t = timesteps.reward[1:]
      # Discount ought to be zero on a LAST timestep, use the mask to ensure this.
      discount_t = timesteps.discount[1:] * mask[1:]
      q_tm1 = q[:-1]
      q_t = q[1:]
      mask_tm1 = mask[:-1]
  
      # Mask out TD errors for the last state in an episode.
      td_error_tm1 = mask_tm1 * rlax.q_lambda(
          q_tm1, a_tm1, r_t, discount_t, q_t, lambda_=self._lambda
      )
      return jnp.sum(rlax.l2_loss(td_error_tm1)) / jnp.sum(mask_tm1)
  
imports: import collections
  from absl import app
  from absl import flags
  from bsuite.environments import catch
  import dm_env
  import haiku
  from haiku import nets
  import jax
  import jax.numpy
  import numpy
  import optax
  import rlax
  from rlax.examples import experiment
label_desc: Use image node_id 7 for calling the OnlineQLambda obj's underlying member method code with example usage: obj._loss(params, actions, timesteps) and returns: unknown