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","."]
docstr: Roll out the dynamics model for each embedding state.
docstr_tok: ["Roll","out","the","dynamics","model","for","each","embedding","state","."]
code: def _compute_model_rollout_predictions(
      self,
      params: mpo_networks.MPONetworkParams,
      state_embeddings: types.NestedArray,
      action_sequence: types.NestedArray,
  ) -> mpo_types.ModelOutputs:
      """Roll out the dynamics model for each embedding state."""
      assert self._model_rollout_length > 0
      # Stack the R=T-K+1 action slices of length K: [0:K; ...; T-K:T]; [K, R].
      rollout_actions = self._rolling_window(action_sequence)
  
      # Create batch of root states (embeddings) s_t for t \in {0, ..., R}.
      num_rollouts = (
          action_sequence.shape[0] - self._model_rollout_length + 1
      )
      root_state = self._dynamics_model.initial_state_fn(
          params.dynamics_model_initial_state,
          state_embeddings[:num_rollouts],
      )
      # TODO(abef): randomly choose (fewer?) root unroll states, as in Muesli?
  
      # Roll out K steps forward in time for each root embedding; [K, R, ...].
      # For example, policy_rollout[k, t] is the step-k prediction starting from
      # state s_t (and same for value_rollout and reward_rollout). Thus, for
      # valid values of k, t, and i, policy_rollout[k, t] and
      # policy_rollout[k-i, t+i] share the same target.
      (
          policy_rollout,
          value_rollout,
          reward_rollout,
          embedding_rollout,
      ), _ = self._dynamics_model.unroll(
          params.dynamics_model, rollout_actions, root_state
      )
      # TODO(abef): try using the same params for both the root & rollout heads.
  
      chex.assert_shape(
          [rollout_actions, embedding_rollout],
          (self._model_rollout_length, num_rollouts, ...),
      )
  
      # Create the outputs but drop the rollout that uses action a_{T-1} (and
      # thus contains state s_T) for the policy, value, and embedding because we
      # don't have targets for s_T (but we do know them for the final reward).
      # Also drop the rollout with s_{T-1} for the value because we don't have
      # targets for that either.
      return mpo_types.ModelOutputs(
          policy=policy_rollout[:, :-1],  # [K, R-1, ...]
          value=value_rollout[:, :-2],  # [K, R-2, ...]
          reward=reward_rollout,  # [K, R, ...]
          embedding=embedding_rollout[:, :-1],
      )
  
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 3 for calling the RolloutLoss obj's underlying member method code with example usage: obj._compute_model_rollout_predictions(params, state_embeddings, action_sequence) and returns: mpo_types