# 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
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




class MPCAgent(object):

  def __init__(self, batch_env, step, is_training, should_log, config):
    self._step = step  # Trainer step, not environment step.
    self._is_training = is_training
    self._should_log = should_log
    self._config = config
    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)

# 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

#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




class MPCAgent(object):

  def __init__(self, batch_env, step, is_training, should_log, config):
    self._step = step  # Trainer step, not environment step.
    self._is_training = is_training
    self._should_log = should_log
    self._config = config
    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)