class_docstr: Model-Based Offline Planning (MBOP) learner.
  
  See https://arxiv.org/abs/2008.05556 for more information.
class_docstr_tok: ["Model-Based","Offline","Planning","(","MBOP",")","learner",".","See","https",":","\/\/arxiv.org\/abs\/2008.05556","for","more","information","."]
docstr: Creates an MBOP learner.
  
  Args:
    networks: One network per model.
    losses: One loss per model.
    iterator: An iterator of time-batched transitions used to train the
      networks.
    rng_key: Random key.
    logger_fn: Constructs a logger for a label.
    make_world_model_learner: Function to create the world model learner.
    make_policy_prior_learner: Function to create the policy prior learner.
    make_n_step_return_learner: Function to create the n-step return learner.
    counter: Parent counter object.
docstr_tok: ["Creates","an","MBOP","learner",".","Args",":","networks",":","One","network","per","model",".","losses",":","One","loss","per","model",".","iterator",":","An","iterator","of","time-batched","transitions","used","to","train","the","networks",".","rng_key",":","Random","key",".","logger_fn",":","Constructs","a","logger","for","a","label",".","make_world_model_learner",":","Function","to","create","the","world","model","learner",".","make_policy_prior_learner",":","Function","to","create","the","policy","prior","learner",".","make_n_step_return_learner",":","Function","to","create","the","n-step","return","learner",".","counter",":","Parent","counter","object","."]
code: def __init__(
      self,
      networks: mbop_networks.MBOPNetworks,
      losses: mbop_losses.MBOPLosses,
      iterator: Iterator[types.Transition],
      rng_key: jax_types.PRNGKey,
      logger_fn: LoggerFn,
      make_world_model_learner: MakeWorldModelLearner,
      make_policy_prior_learner: MakePolicyPriorLearner,
      make_n_step_return_learner: MakeNStepReturnLearner,
      counter: Optional[counting.Counter] = None,
  ):
      """Creates an MBOP learner.
  
      Args:
        networks: One network per model.
        losses: One loss per model.
        iterator: An iterator of time-batched transitions used to train the
          networks.
        rng_key: Random key.
        logger_fn: Constructs a logger for a label.
        make_world_model_learner: Function to create the world model learner.
        make_policy_prior_learner: Function to create the policy prior learner.
        make_n_step_return_learner: Function to create the n-step return learner.
        counter: Parent counter object.
      """
      # General learner book-keeping and loggers.
      self._counter = counter or counting.Counter()
      self._logger = logger_fn("mbop", "steps")
  
      # Prepare iterators for the learners, to not split the data (preserve sample
      # efficiency).
      sharded_prefetching_dataset = utils.sharded_prefetch(iterator)
      (
          world_model_iterator,
          policy_prior_iterator,
          n_step_return_iterator,
      ) = itertools.tee(sharded_prefetching_dataset, 3)
  
      (
          world_model_key,
          policy_prior_key,
          n_step_return_key,
      ) = jax.random.split(rng_key, 3)
  
      self._world_model = make_world_model_learner(
          logger_fn,
          self._counter,
          world_model_key,
          world_model_iterator,
          networks.world_model_network,
          losses.world_model_loss,
      )
      self._policy_prior = make_policy_prior_learner(
          logger_fn,
          self._counter,
          policy_prior_key,
          policy_prior_iterator,
          networks.policy_prior_network,
          losses.policy_prior_loss,
      )
      self._n_step_return = make_n_step_return_learner(
          logger_fn,
          self._counter,
          n_step_return_key,
          n_step_return_iterator,
          networks.n_step_return_network,
          losses.n_step_return_loss,
      )
      # Start recording timestamps after the first learning step to not report
      # "warmup" time.
      self._timestamp = None
      self._learners = {
          "world_model": self._world_model,
          "policy_prior": self._policy_prior,
          "n_step_return": self._n_step_return,
      }
  
imports: import dataclasses
  import functools
  import itertools
  import time
  from typing import Any, Callable, Iterator, List, Optional
  from acme import core
  from acme import types
  from acme.agents.jax import bc
  from acme.agents.jax.mbop import ensemble
  from acme.agents.jax.mbop import losses
  from acme.agents.jax.mbop import networks
  from acme.jax import networks
  from acme.jax import types
  from acme.jax import utils
  from acme.utils import counting
  from acme.utils import loggers
  import jax
  import jax.numpy
  import optax
label_desc: Use image node_id 1 to create a new MBOPLearner object from inherited base classes: core with example: obj = MBOPLearner(networks, losses, iterator, rng_key, logger_fn, make_world_model_learner, make_policy_prior_learner, make_n_step_return_learner, counter)