class_docstr: A MuZero/Muesli-style loss on the rollouts of the dynamics model.
class_docstr_tok: ["A","MuZero\/Muesli-style","loss","on","the","rollouts","of","the","dynamics","model","."]
code: def __call__(
      self,
      params: mpo_networks.MPONetworkParams,
      dual_params: mpo_types.DualParams,
      sequence: adders.Step,
      state_embeddings: types.NestedArray,
      targets: mpo_types.LossTargets,
      key: network_lib.PRNGKey,
  ) -> Tuple[jnp.ndarray, mpo_types.LogDict]:
      num_rollouts = (
          sequence.reward.shape[0] - self._model_rollout_length + 1
      )
      indices = jnp.arange(num_rollouts)
  
      # Create rollout predictions.
      rollout = self._compute_model_rollout_predictions(
          params=params,
          state_embeddings=state_embeddings,
          action_sequence=sequence.action,
      )
  
      # Create rollout target tensors. The rollouts will not contain the policy
      # and value at t=0 because they start after taking the first action in
      # the sequence, so drop those when creating the targets. They will contain
      # the reward at t=0, however, because of how the sequences are stored.
      # Rollout target shapes:
      #   - value: [N, Z, T-2] -> [N, Z, K, R-2],
      #   - reward: [T] -> [K, R].
      value_targets = self._rolling_window(
          targets.value[..., 1:], axis=-1
      )
      reward_targets = self._rolling_window(targets.reward)[
          None, None, ...
      ]
  
      # Define the value and reward rollout loss functions.
      def value_loss_fn(root_idx) -> jnp.ndarray:
          return self._distributional_loss_fn(
              rollout.value[:, root_idx],  # [K, R-2, ...]
              value_targets[..., root_idx],
          )  # [..., K, R-2]
  
      def reward_loss_fn(root_idx) -> jnp.ndarray:
          return self._distributional_loss_fn(
              rollout.reward[:, root_idx],  # [K, R, ...]
              reward_targets[..., root_idx],
          )  # [..., K, R]
  
      # Reward and value losses.
      critic_loss = jnp.mean(jax.vmap(value_loss_fn)(indices[:-2]))
      reward_loss = jnp.mean(jax.vmap(reward_loss_fn)(indices))
  
      # Define the MPO policy rollout loss.
      mpo_policy_loss = 0
      if self._loss_scales.rollout.policy:
          # Rollout target shapes:
          #   - policy: [T-1, ...] -> [K, R-1, ...],
          #   - q_improvement: [N, T-1] -> [N, K, R-1].
          policy_targets = self._rolling_window(targets.policy[1:])
          q_improvement = self._rolling_window(
              targets.q_improvement[:, 1:], axis=1
          )
  
          def policy_loss_fn(root_idx) -> jnp.ndarray:
              chex.assert_shape(
                  (
                      rollout.policy.logits,
                      policy_targets.logits,
                  ),  # pytype: disable=attribute-error  # numpy-scalars
                  (self._model_rollout_length, num_rollouts - 1, None),
              )
              chex.assert_shape(
                  q_improvement,
                  (None, self._model_rollout_length, num_rollouts - 1),
              )
              # Compute MPO's E-step unnormalized logits.
              temperature = discrete_losses.get_temperature_from_params(
                  dual_params
              )
              policy_target_probs = jax.nn.softmax(
                  jnp.transpose(q_improvement[..., root_idx])
                  / temperature
                  + jax.nn.log_softmax(
                      policy_targets[:, root_idx].logits, axis=-1
                  )
              )  # pytype: disable=attribute-error  # numpy-scalars
              return softmax_cross_entropy(
                  rollout.policy[
                      :, root_idx
                  ].logits,  # pytype: disable=bad-return-type  # numpy-scalars
                  jax.lax.stop_gradient(policy_target_probs),
              )
  
          # Compute the MPO loss and add it to the overall rollout policy loss.
          mpo_policy_loss = jax.vmap(policy_loss_fn)(indices[:-1])
          mpo_policy_loss = jnp.mean(mpo_policy_loss)
  
      # Define the BC policy rollout loss (only supported for discrete policies).
      bc_policy_loss, bc_policy_acc = 0, 0
      if self._loss_scales.rollout.bc_policy:
          num_actions = rollout.policy.logits.shape[-1]  # A
          bc_targets = self._rolling_window(  # [T-1, A] -> [K, R-1, A]
              rlax.one_hot(sequence.action[1:], num_actions)
          )
  
          def bc_policy_loss_fn(
              root_idx,
          ) -> Tuple[jnp.ndarray, jnp.ndarray]:
              """Self-behavior-cloning loss (cross entropy on rollout actions)."""
              chex.assert_shape(
                  (rollout.policy.logits, bc_targets),
                  (
                      self._model_rollout_length,
                      num_rollouts - 1,
                      num_actions,
                  ),
              )
              loss = softmax_cross_entropy(
                  rollout.policy.logits[:, root_idx],
                  bc_targets[:, root_idx],
              )
              top1_accuracy = top1_accuracy_tiebreak(
                  rollout.policy.logits[:, root_idx],
                  bc_targets[:, root_idx],
                  rng=key,
              )
              return (
                  loss,
                  top1_accuracy,
              )  # pytype: disable=bad-return-type  # numpy-scalars
  
          # Compute each rollout loss by vmapping over the rollouts.
          bc_policy_loss, bc_policy_acc = jax.vmap(bc_policy_loss_fn)(
              indices[:-1]
          )
          bc_policy_loss = jnp.mean(bc_policy_loss)
          bc_policy_acc = jnp.mean(bc_policy_acc)
  
      # Combine losses.
      loss = (
          self._loss_scales.rollout.policy * mpo_policy_loss
          + self._loss_scales.rollout.bc_policy * bc_policy_loss
          + self._loss_scales.critic
          * self._loss_scales.rollout.critic
          * critic_loss
          + self._loss_scales.rollout.reward * reward_loss
      )
  
      logging_dict = {
          "rollout_critic_loss": critic_loss,
          "rollout_reward_loss": reward_loss,
          "rollout_policy_loss": mpo_policy_loss,
          "rollout_bc_policy_loss": bc_policy_loss,
          "rollout_bc_accuracy": bc_policy_acc,
          "rollout_loss": loss,
      }
  
      return loss, logging_dict
  
file_path: acme/acme/agents/jax/mpo/rollout_loss.py
imports: from typing import Tuple
  from acme import types
  from acme.adders import reverb
  from acme.agents.jax.mpo import categorical_mpo
  from acme.agents.jax.mpo import networks
  from acme.agents.jax.mpo import types
  from acme.agents.jax.mpo import utils
  from acme.jax import networks
  import chex
  import jax
  import jax.numpy
  import rlax
label_desc: Use image node_id 4 for calling the RolloutLoss obj's underlying member method code with example usage: obj.__call__(params, dual_params, sequence, state_embeddings, targets, key) and returns: loss,  logging_dict, self, self, softmax_cross_entropy, loss,  top1_accuracy