class_docstr: Recurrent CRR learner.
  
  This is the learning component of a RCRR agent. It takes a dataset as
  input and implements update functionality to learn from this dataset.
class_docstr_tok: ["Recurrent","CRR","learner",".","This","is","the","learning","component","of","a","RCRR","agent",".","It","takes","a","dataset","as","input","and","implements","update","functionality","to","learn","from","this","dataset","."]
code: def _step(self, sample: reverb.ReplaySample) -> Dict[str, tf.Tensor]:
      # Transpose batch and sequence axes, i.e. [B, T, ...] to [T, B, ...].
      sample = tf2_utils.batch_to_sequence(sample)
      observations = sample.observation
      actions = sample.action
      rewards = sample.reward
      discounts = sample.discount
  
      dtype = rewards.dtype
  
      # Cast the additional discount to match the environment discount dtype.
      discount = tf.cast(self._discount, dtype=discounts.dtype)
  
      # Loss cumulants across time. These cannot be python mutable objects.
      critic_loss = 0.0
      policy_loss = 0.0
  
      # Each transition induces a policy loss, which we then weight using
      # the `policy_loss_coef_t`; shape [B], see https://arxiv.org/abs/2006.15134.
      # `policy_loss_coef` is a scalar average of these coefficients across
      # the batch and sequence length dimensions.
      policy_loss_coef = 0.0
  
      per_device_batch_size = actions.shape[1]
  
      # Initialize recurrent states.
      critic_state = self._critic_network.initial_state(
          per_device_batch_size
      )
      target_critic_state = critic_state
      policy_state = self._policy_network.initial_state(
          per_device_batch_size
      )
      target_policy_state = policy_state
  
      with tf.GradientTape(persistent=True) as tape:
          for t in range(1, self._sequence_length):
              o_tm1 = tree.map_structure(
                  operator.itemgetter(t - 1), observations
              )
              a_tm1 = tree.map_structure(
                  operator.itemgetter(t - 1), actions
              )
              r_t = tree.map_structure(
                  operator.itemgetter(t - 1), rewards
              )
              d_t = tree.map_structure(
                  operator.itemgetter(t - 1), discounts
              )
              o_t = tree.map_structure(
                  operator.itemgetter(t), observations
              )
  
              if t != 1:
                  # By only updating the target critic state here we are forcing
                  # the target critic to ignore observations[0]. Otherwise, the
                  # target_critic will be unrolled for one more timestep than critic.
                  # The smaller the sequence length, the more problematic this is: if
                  # you use RNN on sequences of length 2, you would expect the code to
                  # never use recurrent connections. But if you don't skip updating the
                  # target_critic_state on observation[0] here, it won't be the case.
                  _, target_critic_state = self._target_critic_network(
                      o_tm1, a_tm1, target_critic_state
                  )
  
              # ========================= Critic learning ============================
              q_tm1, next_critic_state = self._critic_network(
                  o_tm1, a_tm1, critic_state
              )
              (
                  target_action_distribution,
                  target_policy_state,
              ) = self._target_policy_network(o_t, target_policy_state)
  
              sampled_actions_t = target_action_distribution.sample(
                  self._num_action_samples_td_learning
              )
              # [N, B, ...]
              tiled_o_t = tf2_utils.tile_nested(
                  o_t, self._num_action_samples_td_learning
              )
              tiled_target_critic_state = tf2_utils.tile_nested(
                  target_critic_state,
                  self._num_action_samples_td_learning,
              )
  
              # Compute the target critic's Q-value of the sampled actions.
              sampled_q_t, _ = snt.BatchApply(
                  self._target_critic_network
              )(tiled_o_t, sampled_actions_t, tiled_target_critic_state)
  
              # Compute average logits by first reshaping them to [N, B, A] and then
              # normalizing them across atoms.
              new_shape = [
                  self._num_action_samples_td_learning,
                  r_t.shape[0],
                  -1,
              ]
              sampled_logits = tf.reshape(sampled_q_t.logits, new_shape)
              sampled_logprobs = tf.math.log_softmax(
                  sampled_logits, axis=-1
              )
              averaged_logits = tf.reduce_logsumexp(
                  sampled_logprobs, axis=0
              )
  
              # Construct the expected distributional value for bootstrapping.
              q_t = networks.DiscreteValuedDistribution(
                  values=sampled_q_t.values, logits=averaged_logits
              )
              critic_loss_t = losses.categorical(
                  q_tm1, r_t, discount * d_t, q_t
              )
              critic_loss_t = tf.reduce_mean(critic_loss_t)
  
              # ========================= Actor learning =============================
              (
                  action_distribution_tm1,
                  policy_state,
              ) = self._policy_network(o_tm1, policy_state)
              q_tm1_mean = q_tm1.mean()
  
              # Compute the estimate of the value function based on
              # self._num_action_samples_policy_weight samples from the policy.
              tiled_o_tm1 = tf2_utils.tile_nested(
                  o_tm1, self._num_action_samples_policy_weight
              )
              tiled_critic_state = tf2_utils.tile_nested(
                  critic_state, self._num_action_samples_policy_weight
              )
              action_tm1 = action_distribution_tm1.sample(
                  self._num_action_samples_policy_weight
              )
              tiled_z_tm1, _ = snt.BatchApply(self._critic_network)(
                  tiled_o_tm1, action_tm1, tiled_critic_state
              )
              tiled_v_tm1 = tf.reshape(
                  tiled_z_tm1.mean(),
                  [self._num_action_samples_policy_weight, -1],
              )
  
              # Use mean, min, or max to aggregate Q(s, a_i), a_i ~ pi(s) into the
              # final estimate of the value function.
              if self._baseline_reduce_function == "mean":
                  v_tm1_estimate = tf.reduce_mean(tiled_v_tm1, axis=0)
              elif self._baseline_reduce_function == "max":
                  v_tm1_estimate = tf.reduce_max(tiled_v_tm1, axis=0)
              elif self._baseline_reduce_function == "min":
                  v_tm1_estimate = tf.reduce_min(tiled_v_tm1, axis=0)
  
              # Assert that action_distribution_tm1 is a batch of multivariate
              # distributions (in contrast to e.g. a [batch, action_size] collection
              # of 1d distributions).
              assert len(action_distribution_tm1.batch_shape) == 1
              policy_loss_batch = -action_distribution_tm1.log_prob(
                  a_tm1
              )
  
              advantage = q_tm1_mean - v_tm1_estimate
              if self._policy_improvement_modes == "exp":
                  policy_loss_coef_t = tf.math.minimum(
                      tf.math.exp(advantage / self._beta),
                      self._ratio_upper_bound,
                  )
              elif self._policy_improvement_modes == "binary":
                  policy_loss_coef_t = tf.cast(
                      advantage > 0, dtype=dtype
                  )
              elif self._policy_improvement_modes == "all":
                  # Regress against all actions (effectively pure BC).
                  policy_loss_coef_t = 1.0
              policy_loss_coef_t = tf.stop_gradient(policy_loss_coef_t)
  
              policy_loss_batch *= policy_loss_coef_t
              policy_loss_t = tf.reduce_mean(policy_loss_batch)
  
              critic_state = next_critic_state
  
              critic_loss += critic_loss_t
              policy_loss += policy_loss_t
              policy_loss_coef += tf.reduce_mean(
                  policy_loss_coef_t
              )  # For logging.
  
          # Divide by sequence length to get mean losses.
          critic_loss /= tf.cast(self._sequence_length, dtype=dtype)
          policy_loss /= tf.cast(self._sequence_length, dtype=dtype)
          policy_loss_coef /= tf.cast(
              self._sequence_length, dtype=dtype
          )
  
      # Compute gradients.
      critic_gradients = tape.gradient(
          critic_loss, self._critic_network.trainable_variables
      )
      policy_gradients = tape.gradient(
          policy_loss, self._policy_network.trainable_variables
      )
  
      # Delete the tape manually because of the persistent=True flag.
      del tape
  
      # Sync gradients across GPUs or TPUs.
      ctx = tf.distribute.get_replica_context()
      critic_gradients = ctx.all_reduce("mean", critic_gradients)
      policy_gradients = ctx.all_reduce("mean", policy_gradients)
  
      # Maybe clip gradients.
      if self._clipping:
          policy_gradients = tf.clip_by_global_norm(
              policy_gradients, 40.0
          )[0]
          critic_gradients = tf.clip_by_global_norm(
              critic_gradients, 40.0
          )[0]
  
      # Apply gradients.
      self._critic_optimizer.apply(
          critic_gradients, self._critic_network.trainable_variables
      )
      self._policy_optimizer.apply(
          policy_gradients, self._policy_network.trainable_variables
      )
  
      source_variables = (
          self._critic_network.variables
          + self._policy_network.variables
      )
      target_variables = (
          self._target_critic_network.variables
          + self._target_policy_network.variables
      )
  
      # Make online -> target network update ops.
      if tf.math.mod(self._num_steps, self._target_update_period) == 0:
          for src, dest in zip(source_variables, target_variables):
              dest.assign(src)
      self._num_steps.assign_add(1)
  
      return {
          "critic_loss": critic_loss,
          "policy_loss": policy_loss,
          "policy_loss_coef": policy_loss_coef,
      }
  
imports: import operator
  import time
  from typing import Dict, List, Optional
  from acme import core
  from acme.tf import losses
  from acme.tf import networks
  from acme.tf import savers
  from acme.tf import utils
  from acme.utils import counting
  from acme.utils import loggers
  import numpy
  import reverb
  import sonnet
  import tensorflow
  import tree
label_desc: Use image node_id 2 for calling the RCRRLearner obj's underlying member method code with example usage: obj._step(sample) and returns: dict