code: def make_world_model_network(
      spec: specs.EnvironmentSpec,
      hidden_layer_sizes: Tuple[int, ...] = (64, 64),
  ) -> networks.FeedForwardNetwork:
      """Creates a world model network used by the agent."""
  
      observation_size = np.prod(spec.observations.shape, dtype=int)
  
      def _world_model_fn(
          observation_t, action_t, is_training=False, key=None
      ):
          # is_training and key allows to defined train/test dependant modules
          # like dropout.
          del is_training
          del key
          network = hk.nets.MLP(
              hidden_layer_sizes + (observation_size + 1,)
          )
          # World model returns both an observation and a reward.
          observation_tp1, reward_t = jnp.split(
              network(
                  jnp.concatenate([observation_t, action_t], axis=-1)
              ),
              [observation_size],
              axis=-1,
          )
          return observation_tp1, reward_t
  
      world_model = hk.without_apply_rng(hk.transform(_world_model_fn))
      return make_network_from_module(world_model, spec)
  
file_path: acme/acme/agents/jax/mbop/networks.py
imports: import dataclasses
  from typing import Any, Tuple
  from acme import specs
  from acme.jax import networks
  from acme.jax import utils
  import haiku
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: make_world_model_network(spec, hidden_layer_sizes) and returns: make_network_from_module, observation_tp1,  reward_t