# Path: dreamer/tools/nested.py
# def zip_(*structures, **kwargs):
# def map_(function, *structures, **kwargs):
#   def impl(function, *structures):
# def flatten_(structure):
# def filter_(predicate, *structures, **kwargs):
#   def impl(predicate, *structures):


import numpy as np
import tensorflow as tf
from tensorflow_probability import distributions as tfd
from dreamer.tools import nested
    self._cell = config.cell
    self._num_envs = len(batch_env)
    state = self._cell.zero_state(self._num_envs, tf.float32)
    var_like = lambda x: tf.get_local_variable(
        x.name.split(':')[0].replace('/', '_') + '_var',
        shape=x.shape,
        initializer=lambda *_, **__: tf.zeros_like(x), use_resource=True)
    self._state = nested.map(var_like, state)
    batch_action_shape = (self._num_envs,) + batch_env.action_space.shape
    self._prev_action = tf.get_local_variable(
        'prev_action_var', shape=batch_action_shape,
        initializer=lambda *_, **__: tf.zeros(batch_action_shape),
        use_resource=True)

  def reset(self, agent_indices):
    state = nested.map(
        lambda tensor: tf.gather(tensor, agent_indices),
        self._state)
    reset_state = nested.map(
        lambda var, val: tf.scatter_update(var, agent_indices, 0 * val),
        self._state, state, flatten=True)
    reset_prev_action = self._prev_action.assign(
        tf.zeros_like(self._prev_action))
    return tf.group(reset_prev_action, *reset_state)

  def step(self, agent_indices, observ):
    observ = self._config.preprocess_fn(observ)
    # Converts observ to sequence.
    observ = nested.map(lambda x: x[:, None], observ)
    embedded = self._config.encoder(observ)[:, 0]

# Path: dreamer/tools/nested.py
# def zip_(*structures, **kwargs):
# def map_(function, *structures, **kwargs):
#   def impl(function, *structures):
# def flatten_(structure):
# def filter_(predicate, *structures, **kwargs):
#   def impl(predicate, *structures):


# Path: dreamer/control/mpc_agent.py
import numpy as np
import tensorflow as tf
from tensorflow_probability import distributions as tfd
from dreamer.tools import nested

    self._cell = config.cell
    self._num_envs = len(batch_env)
    state = self._cell.zero_state(self._num_envs, tf.float32)
    var_like = lambda x: tf.get_local_variable(
        x.name.split(':')[0].replace('/', '_') + '_var',
        shape=x.shape,
        initializer=lambda *_, **__: tf.zeros_like(x), use_resource=True)
    self._state = nested.map(var_like, state)
    batch_action_shape = (self._num_envs,) + batch_env.action_space.shape
    self._prev_action = tf.get_local_variable(
        'prev_action_var', shape=batch_action_shape,
        initializer=lambda *_, **__: tf.zeros(batch_action_shape),
        use_resource=True)

  def reset(self, agent_indices):
    state = nested.map(
        lambda tensor: tf.gather(tensor, agent_indices),
        self._state)
    reset_state = nested.map(
        lambda var, val: tf.scatter_update(var, agent_indices, 0 * val),
        self._state, state, flatten=True)
    reset_prev_action = self._prev_action.assign(
        tf.zeros_like(self._prev_action))
    return tf.group(reset_prev_action, *reset_state)

  def step(self, agent_indices, observ):
    observ = self._config.preprocess_fn(observ)
    # Converts observ to sequence.
    observ = nested.map(lambda x: x[:, None], observ)
    embedded = self._config.encoder(observ)[:, 0]

    state = nested.map(