class_docstr: MPO learner (discrete or continuous, distributional or not).
class_docstr_tok: ["MPO","learner","(","discrete","or","continuous",",","distributional","or","not",")","."]
docstr: Compute the targets needed to train the agent.
docstr_tok: ["Compute","the","targets","needed","to","train","the","agent","."]
code: def _compute_targets(
      self,
      target_params: mpo_networks.MPONetworkParams,
      dual_params: mpo_types.DualParams,
      sequence: adders.Step,
      online_policy: types.NestedArray,  # TODO(abef): remove this.
      key: jax_types.PRNGKey,
  ) -> mpo_types.LossTargets:
      """Compute the targets needed to train the agent."""
  
      # 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(
              target_params.torso_initial_state, None
          )
  
      # Unroll the target core network. Note that this may pass the embeddings
      # unchanged if, say, the core is an hk.IdentityCore.
      target_state_embedding, _ = self._networks.torso_unroll(
          target_params, sequence.observation, initial_state
      )  # [T, ...]
  
      # Compute the action distribution from target policy network.
      target_policy = self._networks.policy_head_apply(
          target_params, target_state_embedding
      )  # [T, ...]
  
      # Maybe reward clip.
      clipped_reward = jnp.clip(
          sequence.reward, *self._reward_clip_range
      )  # [T]
      # TODO(abef): when to clip rewards, if at all, if learning dynamics model?
  
      @jax.named_call
      @jax.vmap
      def critic_mean_fn(action_: jnp.ndarray) -> jnp.ndarray:
          """Compute mean of target critic distribution."""
          critic_output = self._networks.critic_head_apply(
              target_params, target_state_embedding, action_
          )
          if self._critic_type != CriticType.NONDISTRIBUTIONAL:
              critic_output = critic_output.mean()
          return critic_output
  
      @jax.named_call
      @jax.vmap
      def critic_sample_fn(
          action_: jnp.ndarray, seed_: jnp.ndarray
      ) -> jnp.ndarray:
          """Sample from the target critic distribution."""
          z_distribution = self._networks.critic_head_apply(
              target_params, target_state_embedding, action_
          )
          z_samples = z_distribution.sample(
              self._policy_eval_num_val_samples, seed=seed_
          )
          return z_samples  # [Z, T, 1]
  
      if self._discrete_policy:
          # Use all actions to improve policy (no sampling); N = num_actions.
          a_improvement = jnp.arange(
              self._action_spec.num_values
          )  # [N]
          seq_len = target_state_embedding.shape[0]  # T
          a_improvement = jnp.tile(
              a_improvement[..., None], [1, seq_len]
          )  # [N, T]
      else:
          # Sample actions to improve policy; [N=num_samples, T].
          a_improvement = target_policy.sample(
              self._num_samples, seed=key
          )
  
      # TODO(abef): use model to get q_improvement = r + gamma*V?
  
      # Compute the mean Q-values used in policy improvement; [N, T].
      q_improvement = critic_mean_fn(a_improvement).squeeze(axis=-1)
  
      # Policy to use for policy evaluation and bootstrapping.
      if self._use_online_policy_to_bootstrap:
          policy_to_evaluate = online_policy
          chex.assert_equal(
              online_policy.batch_shape, target_policy.batch_shape
          )
      else:
          policy_to_evaluate = target_policy
  
      # Action(s) to use for policy evaluation; shape [N, T].
      if self._policy_eval_stochastic:
          a_evaluation = policy_to_evaluate.sample(
              self._num_samples, seed=key
          )
      else:
          a_evaluation = policy_to_evaluate.mode()
          a_evaluation = jnp.expand_dims(
              a_evaluation, axis=0
          )  # [N=1, T]
  
      # TODO(abef): policy_eval_stochastic=False makes our targets more "greedy"
  
      # Add a stopgrad in case we use the online policy for evaluation.
      a_evaluation = jax.lax.stop_gradient(a_evaluation)
  
      if self._critic_type == CriticType.MIXTURE_OF_GAUSSIANS:
          # Produce Z return samples for every N action sample; [N, Z, T, 1].
          seeds = jax.random.split(key, num=a_evaluation.shape[0])
          z_samples = critic_sample_fn(a_evaluation, seeds)
      else:
          normalized_weights = 1.0 / a_evaluation.shape[0]
          z_samples = critic_mean_fn(a_evaluation)  # [N, T, 1]
  
          # When policy_eval_stochastic == True, this corresponds to expected SARSA.
          # Otherwise, normalized_weights = 1.0 and N = 1 so the sum is a no-op.
          z_samples = jnp.sum(
              normalized_weights * z_samples, axis=0, keepdims=True
          )
          z_samples = jnp.expand_dims(
              z_samples, axis=1
          )  # [N, Z=1, T, 1]
  
      # Slice to t = 1...T and transform into raw reward space; [N, Z, T].
      z_samples_itx = self._tx_pair.apply_inv(
          z_samples.squeeze(axis=-1)
      )
  
      # Compute the value estimate by averaging the sampled returns in the raw
      # reward space; shape [N=1, Z=1, T].
      value_target_itx = jnp.mean(
          z_samples_itx, axis=(0, 1), keepdims=True
      )
  
      if self._use_retrace:
          # Warning! Retrace has not been tested with the MoG critic.
          log_rhos = (
              target_policy.log_prob(sequence.action)
              - sequence.extras["log_prob"]
          )
  
          # Compute Q-values; expand and squeeze because critic_mean_fn is vmapped.
          q_t = critic_mean_fn(
              jnp.expand_dims(sequence.action, axis=0)
          ).squeeze(0)
          q_t = q_t.squeeze(
              -1
          )  # Also squeeze trailing scalar dimension; [T].
  
          # Compute retrace targets.
          # These targets use the rewards and discounts as in normal TD-learning but
          # they use a mix of bootstrapped values V(s') and Q(s', a'), weighing the
          # latter based on how likely a' is under the current policy (s' and a' are
          # samples from replay).
          # See [Munos et al., 2016](https://arxiv.org/abs/1606.02647) for more.
          q_value_target_itx = (
              rlax.general_off_policy_returns_from_q_and_v(
                  q_t=self._tx_pair.apply_inv(q_t[1:-1]),
                  v_t=jnp.squeeze(value_target_itx, axis=(0, 1))[1:],
                  r_t=clipped_reward[:-1],
                  discount_t=self._discount * sequence.discount[:-1],
                  c_t=self._retrace_lambda
                  * jnp.minimum(1.0, jnp.exp(log_rhos[1:-1])),
              )
          )
  
          # Expand dims to the expected [N=1, Z=1, T-1].
          q_value_target_itx = jnp.expand_dims(
              q_value_target_itx, axis=(0, 1)
          )
      else:
          # Compute bootstrap target from sequences. vmap return computation across
          # N action and Z return samples; shape [N, Z, T-1].
          n_step_return_fn = functools.partial(
              rlax.n_step_bootstrapped_returns,
              r_t=clipped_reward[:-1],
              discount_t=self._discount * sequence.discount[:-1],
              n=self._n_step_for_sequence_bootstrap,
              lambda_t=self._td_lambda,
          )
          n_step_return_vfn = jax.vmap(jax.vmap(n_step_return_fn))
          q_value_target_itx = n_step_return_vfn(
              v_t=z_samples_itx[..., 1:]
          )
  
      # Transform back to the canonical space and stop gradients.
      q_value_target = jax.lax.stop_gradient(
          self._tx_pair.apply(q_value_target_itx)
      )
      reward_target = jax.lax.stop_gradient(
          self._tx_pair.apply(clipped_reward)
      )
      value_target = jax.lax.stop_gradient(
          self._tx_pair.apply(value_target_itx)
      )
  
      if self._critic_type == mpo_types.CriticType.CATEGORICAL:
  
          @jax.vmap
          def get_logits_and_values(
              action: jnp.ndarray,
          ) -> Tuple[jnp.ndarray, jnp.ndarray]:
              critic_output = self._networks.critic_head_apply(
                  target_params, target_state_embedding[1:], action
              )
              return critic_output.logits, critic_output.values
  
          z_t_logits, z_t_values = get_logits_and_values(
              a_evaluation[:, 1:]
          )
          z_t_logits = jnp.squeeze(z_t_logits, axis=2)  # [N, T-1, L]
          z_t_values = z_t_values[
              0
          ]  # Values are identical at each N; [L].
  
          gamma = (
              self._discount * sequence.discount[:-1, None]
          )  # [T-1, 1]
          r_t = clipped_reward[:-1, None]  # [T-1, 1]
          atoms_itx = self._tx_pair.apply_inv(z_t_values)[
              None, ...
          ]  # [1, L]
          z_target_atoms = self._tx_pair.apply(
              r_t + gamma * atoms_itx
          )  # [T-1, L]
          # Note: this is n=1-step TD unless using experience=FromTransitions(n>1).
          z_target_probs = jax.nn.softmax(z_t_logits)  # [N, T-1, L]
          z_target_atoms = jax.lax.broadcast(
              z_target_atoms, z_target_probs.shape[:1]
          )  # [N, T-1, L]
          project_fn = functools.partial(
              rlax.categorical_l2_project, z_q=z_t_values
          )
          z_target = jax.vmap(jax.vmap(project_fn))(
              z_target_atoms, z_target_probs
          )
          z_target = jnp.mean(z_target, axis=0)
          q_value_target = jax.lax.stop_gradient(
              z_target[None, ...]
          )  # [1, T-1, L]
          # TODO(abef): make q_v_target shape align with expected [N, Z, T-1] shape?
  
      targets = mpo_types.LossTargets(
          policy=target_policy,  # [T, ...]
          a_improvement=a_improvement,  # [N, T]
          q_improvement=q_improvement,  # [N, T]
          q_value=q_value_target,  # [N, Z, T-1] ([1, T-1, L] for CATEGORICAL)
          value=value_target[..., :-1],  # [N=1, Z=1, T-1]
          reward=reward_target,  # [T]
          embedding=target_state_embedding,
      )  # [T, ...]
  
      return targets
  
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 4 for calling the MPOLearner obj's underlying member method code with example usage: obj._compute_targets(target_params, dual_params, sequence, online_policy, key) and returns: targets, critic_output, z_samples, critic_output,  critic_output