class_docstr: MPO learner (discrete or continuous, distributional or not).
class_docstr_tok: ["MPO","learner","(","discrete","or","continuous",",","distributional","or","not",")","."]
docstr: Compute model predictions at observed and rolled out states.
docstr_tok: ["Compute","model","predictions","at","observed","and","rolled","out","states","."]
code: def _compute_predictions(
      self, params: mpo_networks.MPONetworkParams, sequence: adders.Step
  ) -> mpo_types.ModelOutputs:
      """Compute model predictions at observed and rolled out states."""
  
      # Initialize the core states, possibly to the recorded stale state.
      if self._use_stale_state:
          initial_state = utils.maybe_recover_lstm_type(
              sequence.extras["core_state"]
          )
          initial_state = tree.map_structure(
              lambda x: x[0], initial_state
          )
      else:
          initial_state = self._networks.torso.initial_state_fn(
              params.torso_initial_state, None
          )
  
      # Unroll the online core network. Note that this may pass the embeddings
      # unchanged if, say, the core is an hk.IdentityCore.
      state_embedding, _ = self._networks.torso_unroll(  # [T, ...]
          params, sequence.observation, initial_state
      )
  
      # Compute the root policy and critic outputs; [T, ...] and [T-1, ...].
      policy = self._networks.policy_head_apply(params, state_embedding)
      q_value = self._networks.critic_head_apply(
          params, state_embedding[:-1], sequence.action[:-1]
      )
  
      return mpo_types.ModelOutputs(
          policy=policy,  # [T, ...]
          q_value=q_value,  # [T-1, ...]
          reward=None,
          embedding=state_embedding,
      )
  
imports: import dataclasses
  import functools
  import time
  from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Sequence, Tuple, Union
  from absl import logging
  import acme
  from acme import specs
  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 rollout_loss
  from acme.agents.jax.mpo import types
  from acme.agents.jax.mpo import utils
  from acme.jax import networks
  from acme.jax import types
  from acme.jax import utils
  import acme.jax.losses.mpo
  from acme.utils import counting
  from acme.utils import loggers
  import chex
  import jax
  import jax.numpy
  import numpy
  import optax
  import reverb
  import rlax
  import tree
label_desc: Use image node_id 3 for calling the MPOLearner obj's underlying member method code with example usage: obj._compute_predictions(params, sequence) and returns: mpo_types