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 __init__(
      self,
      dynamics_model: mpo_networks.UnrollableNetwork,
      model_rollout_length: int,
      loss_scales: mpo_types.LossScalesConfig,
      distributional_loss_fn: mpo_types.DistributionalLossFn,
  ):
      self._dynamics_model = dynamics_model
      self._model_rollout_length = model_rollout_length
      self._loss_scales = loss_scales
      self._distributional_loss_fn = distributional_loss_fn
  
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 1 to create a new RolloutLoss object  with example: obj = RolloutLoss(dynamics_model, model_rollout_length, loss_scales, distributional_loss_fn)