Add Dreamerv2 Files

DreamerV2/configs.yaml

@@ -0,0 +1,185 @@
+defaults:
+  gpu: 'none'
+  logdir: ./
+  traindir: null
+  evaldir: null
+  offline_traindir: ''
+  offline_evaldir: ''
+  seed: 0
+  steps: 1e7
+  eval_every: 1e4
+  log_every: 1e4
+  reset_every: 0
+  gpu_growth: True
+  precision: 32
+  debug: False
+  expl_gifs: False
+
+  # Environment
+  task: 'dmc_walker_walk'
+  size: [64, 64]
+  envs: 1
+  action_repeat: 2
+  time_limit: 1000
+  prefill: 2500
+  eval_noise: 0.0
+  clip_rewards: 'identity'
+  atari_grayscale: False
+
+  # Model
+  dyn_cell: 'gru'
+  dyn_hidden: 200
+  dyn_deter: 200
+  dyn_stoch: 50
+  dyn_discrete: 0
+  dyn_input_layers: 1
+  dyn_output_layers: 1
+  dyn_shared: False
+  dyn_mean_act: 'none'
+  dyn_std_act: 'sigmoid2'
+  dyn_min_std: 0.1
+  grad_heads: ['image', 'reward']
+  units: 400
+  reward_layers: 2
+  discount_layers: 3
+  value_layers: 3
+  actor_layers: 4
+  act: 'elu'
+  cnn_depth: 32
+  encoder_kernels: [4, 4, 4, 4]
+  decoder_kernels: [5, 5, 6, 6]
+  decoder_thin: True
+  value_head: 'normal'
+  kl_scale: '1.0'
+  kl_balance: '0.8'
+  kl_free: '1.0'
+  pred_discount: False
+  discount_scale: 1.0
+  reward_scale: 1.0
+  weight_decay: 0.0
+
+  # Training
+  batch_size: 50
+  batch_length: 50
+  train_every: 5
+  train_steps: 1
+  pretrain: 100
+  model_lr: 3e-4
+  value_lr: 8e-5
+  actor_lr: 8e-5
+  opt_eps: 1e-5
+  grad_clip: 100
+  value_grad_clip: 100
+  actor_grad_clip: 100
+  dataset_size: 0
+  oversample_ends: False
+  slow_value_target: True
+  slow_actor_target: True
+  slow_target_update: 100
+  slow_target_fraction: 1
+  opt: 'adam'
+
+  # Behavior.
+  discount: 0.99
+  discount_lambda: 0.95
+  imag_horizon: 15
+  imag_gradient: 'dynamics'
+  imag_gradient_mix: '0.1'
+  imag_sample: True
+  actor_dist: 'trunc_normal'
+  actor_entropy: '1e-4'
+  actor_state_entropy: 0.0
+  actor_init_std: 1.0
+  actor_min_std: 0.1
+  actor_disc: 5
+  actor_temp: 0.1
+  actor_outscale: 0.0
+  expl_amount: 0.0
+  eval_state_mean: False
+  collect_dyn_sample: True
+  behavior_stop_grad: True
+  value_decay: 0.0
+  future_entropy: False
+
+  # Exploration
+  expl_behavior: 'greedy'
+  expl_until: 0
+  expl_extr_scale: 0.0
+  expl_intr_scale: 1.0
+  disag_target: 'stoch'
+  disag_log: True
+  disag_models: 10
+  disag_offset: 1
+  disag_layers: 4
+  disag_units: 400
+
+atari:
+
+  # General
+  task: 'atari_demon_attack'
+  steps: 3e7
+  eval_every: 1e5
+  log_every: 1e4
+  prefill: 50000
+  dataset_size: 2e6
+  pretrain: 0
+  precision: 16
+
+  # Environment
+  time_limit: 108000  # 30 minutes of game play.
+  atari_grayscale: True
+  action_repeat: 4
+  eval_noise: 0.001
+  train_every: 16
+  train_steps: 1
+  clip_rewards: 'tanh'
+
+  # Model
+  grad_heads: ['image', 'reward', 'discount']
+  dyn_cell: 'gru_layer_norm'
+  pred_discount: True
+  cnn_depth: 48
+  dyn_deter: 600
+  dyn_hidden: 600
+  dyn_stoch: 32
+  dyn_discrete: 32
+  reward_layers: 4
+  discount_layers: 4
+  value_layers: 4
+  actor_layers: 4
+
+  # Behavior
+  actor_dist: 'onehot'
+  actor_entropy: 'linear(3e-3,3e-4,2.5e6)'
+  expl_amount: 0.0
+  expl_until: 3e7
+  discount: 0.995
+  imag_gradient: 'both'
+  imag_gradient_mix: 'linear(0.1,0,2.5e6)'
+
+  # Training
+  discount_scale: 5.0
+  reward_scale: 1
+  weight_decay: 1e-6
+  model_lr: 2e-4
+  kl_scale: 0.1
+  kl_free: 0.0
+  actor_lr: 4e-5
+  value_lr: 1e-4
+  oversample_ends: True
+
+  # Disen
+  disen_cnn_depth: 16
+  disen_only_scale: 1.0
+  disen_discount_scale: 2000.0
+  disen_reward_scale: 2000.0
+  num_reward_opt_iters: 20
+
+debug:
+
+  debug: True
+  pretrain: 1
+  prefill: 1
+  train_steps: 1
+  batch_size: 10
+  batch_length: 20

DreamerV2/dreamer.py

@@ -0,0 +1,316 @@
+import argparse
+import collections
+import functools
+import os
+import pathlib
+import sys
+import warnings
+
+warnings.filterwarnings('ignore', '.*box bound precision lowered.*')
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+os.environ['MUJOCO_GL'] = 'egl'
+
+import numpy as np
+import ruamel.yaml as yaml
+import tensorflow as tf
+from tensorflow.keras.mixed_precision import experimental as prec
+
+tf.get_logger().setLevel('ERROR')
+
+from tensorflow_probability import distributions as tfd
+
+sys.path.append(str(pathlib.Path(__file__).parent))
+
+import exploration as expl
+import models
+import tools
+import wrappers
+
+class Dreamer(tools.Module):
+
+  def __init__(self, config, logger, dataset):
+    self._config = config
+    self._logger = logger
+    self._float = prec.global_policy().compute_dtype
+    self._should_log = tools.Every(config.log_every)
+    self._should_train = tools.Every(config.train_every)
+    self._should_pretrain = tools.Once()
+    self._should_reset = tools.Every(config.reset_every)
+    self._should_expl = tools.Until(int(
+        config.expl_until / config.action_repeat))
+    self._metrics = collections.defaultdict(tf.metrics.Mean)
+    with tf.device('cpu:0'):
+      self._step = tf.Variable(count_steps(config.traindir), dtype=tf.int64)
+    # Schedules.
+    config.actor_entropy = (
+        lambda x=config.actor_entropy: tools.schedule(x, self._step))
+    config.actor_state_entropy = (
+        lambda x=config.actor_state_entropy: tools.schedule(x, self._step))
+    config.imag_gradient_mix = (
+        lambda x=config.imag_gradient_mix: tools.schedule(x, self._step))
+    self._dataset = iter(dataset)
+    self._wm = models.WorldModel(self._step, config)
+    self._task_behavior = models.ImagBehavior(
+        config, self._wm, config.behavior_stop_grad)
+    reward = lambda f, s, a: self._wm.heads['reward'](f).mode()
+    self._expl_behavior = dict(
+        greedy=lambda: self._task_behavior,
+        random=lambda: expl.Random(config),
+        plan2explore=lambda: expl.Plan2Explore(config, self._wm, reward),
+    )[config.expl_behavior]()
+    # Train step to initialize variables including optimizer statistics.
+    self._train(next(self._dataset))
+
+  def __call__(self, obs, reset, state=None, training=True):
+    step = self._step.numpy().item()
+    if self._should_reset(step):
+      state = None
+    if state is not None and reset.any():
+      mask = tf.cast(1 - reset, self._float)[:, None]
+      state = tf.nest.map_structure(lambda x: x * mask, state)
+    if training and self._should_train(step):
+      steps = (
+          self._config.pretrain if self._should_pretrain()
+          else self._config.train_steps)
+      for _ in range(steps):
+        self._train(next(self._dataset))
+      if self._should_log(step):
+        for name, mean in self._metrics.items():
+          self._logger.scalar(name, float(mean.result()))
+          mean.reset_states()
+        openl_joint, openl_main, openl_disen, openl_mask = self._wm.video_pred(next(self._dataset))
+        self._logger.video('train_openl_joint', openl_joint)
+        self._logger.video('train_openl_main', openl_main)
+        self._logger.video('train_openl_disen', openl_disen)
+        self._logger.video('train_openl_mask', openl_mask)
+        self._logger.write(fps=True)
+    action, state = self._policy(obs, state, training)
+    if training:
+      self._step.assign_add(len(reset))
+      self._logger.step = self._config.action_repeat \
+          * self._step.numpy().item()
+    return action, state
+
+  @tf.function
+  def _policy(self, obs, state, training):
+    if state is None:
+      batch_size = len(obs['image'])
+      latent = self._wm.dynamics.initial(len(obs['image']))
+      action = tf.zeros((batch_size, self._config.num_actions), self._float)
+    else:
+      latent, action = state
+    embed = self._wm.encoder(self._wm.preprocess(obs))
+    latent, _ = self._wm.dynamics.obs_step(
+        latent, action, embed, self._config.collect_dyn_sample)
+    if self._config.eval_state_mean:
+      latent['stoch'] = latent['mean']
+    feat = self._wm.dynamics.get_feat(latent)
+    if not training:
+      action = self._task_behavior.actor(feat).mode()
+    elif self._should_expl(self._step):
+      action = self._expl_behavior.actor(feat).sample()
+    else:
+      action = self._task_behavior.actor(feat).sample()
+    if self._config.actor_dist == 'onehot_gumble':
+      action = tf.cast(
+          tf.one_hot(tf.argmax(action, axis=-1), self._config.num_actions),
+          action.dtype)
+    action = self._exploration(action, training)
+    state = (latent, action)
+    return action, state
+
+  def _exploration(self, action, training):
+    amount = self._config.expl_amount if training else self._config.eval_noise
+    if amount == 0:
+      return action
+    amount = tf.cast(amount, self._float)
+    if 'onehot' in self._config.actor_dist:
+      probs = amount / self._config.num_actions + (1 - amount) * action
+      return tools.OneHotDist(probs=probs).sample()
+    else:
+      return tf.clip_by_value(tfd.Normal(action, amount).sample(), -1, 1)
+    raise NotImplementedError(self._config.action_noise)
+
+  @tf.function
+  def _train(self, data):
+    print('Tracing train function.')
+    metrics = {}
+    embed, post, feat, kl, mets = self._wm.train(data)
+    metrics.update(mets)
+    start = post
+    if self._config.pred_discount:  # Last step could be terminal.
+      start = {k: v[:, :-1] for k, v in post.items()}
+      embed, feat, kl = embed[:, :-1], feat[:, :-1], kl[:, :-1]
+    reward = lambda f, s, a: self._wm.heads['reward'](f).mode()
+    metrics.update(self._task_behavior.train(start, reward)[-1])
+    if self._config.expl_behavior != 'greedy':
+      mets = self._expl_behavior.train(start, feat, embed, kl)[-1]
+      metrics.update({'expl_' + key: value for key, value in mets.items()})
+    for name, value in metrics.items():
+      self._metrics[name].update_state(value)
+
+
+def count_steps(folder):
+  return sum(int(str(n).split('-')[-1][:-4]) - 1 for n in folder.glob('*.npz'))
+
+
+def make_dataset(episodes, config):
+  example = episodes[next(iter(episodes.keys()))]
+  types = {k: v.dtype for k, v in example.items()}
+  shapes = {k: (None,) + v.shape[1:] for k, v in example.items()}
+  generator = lambda: tools.sample_episodes(
+      episodes, config.batch_length, config.oversample_ends)
+  dataset = tf.data.Dataset.from_generator(generator, types, shapes)
+  dataset = dataset.batch(config.batch_size, drop_remainder=True)
+  dataset = dataset.prefetch(10)
+  return dataset
+
+
+def make_env(config, logger, mode, train_eps, eval_eps):
+  suite, task = config.task.split('_', 1)
+  if suite == 'dmc':
+    env = wrappers.DeepMindControl(task, config.action_repeat, config.size)
+    env = wrappers.NormalizeActions(env)
+  elif suite == 'atari':
+    env = wrappers.Atari(
+        task, config.action_repeat, config.size,
+        grayscale=config.atari_grayscale,
+        life_done=False and (mode == 'train'),
+        sticky_actions=True,
+        all_actions=True)
+    env = wrappers.OneHotAction(env)
+  else:
+    raise NotImplementedError(suite)
+  env = wrappers.TimeLimit(env, config.time_limit)
+  callbacks = [functools.partial(
+      process_episode, config, logger, mode, train_eps, eval_eps)]
+  env = wrappers.CollectDataset(env, callbacks)
+  env = wrappers.RewardObs(env)
+  return env
+
+
+def process_episode(config, logger, mode, train_eps, eval_eps, episode):
+  directory = dict(train=config.traindir, eval=config.evaldir)[mode]
+  cache = dict(train=train_eps, eval=eval_eps)[mode]
+  filename = tools.save_episodes(directory, [episode])[0]
+  length = len(episode['reward']) - 1
+  score = float(episode['reward'].astype(np.float64).sum())
+  video = episode['image']
+  if mode == 'eval':
+    cache.clear()
+  if mode == 'train' and config.dataset_size:
+    total = 0
+    for key, ep in reversed(sorted(cache.items(), key=lambda x: x[0])):
+      if total <= config.dataset_size - length:
+        total += len(ep['reward']) - 1
+      else:
+        del cache[key]
+    logger.scalar('dataset_size', total + length)
+  cache[str(filename)] = episode
+  print(f'{mode.title()} episode has {length} steps and return {score:.1f}.')
+  logger.scalar(f'{mode}_return', score)
+  logger.scalar(f'{mode}_length', length)
+  logger.scalar(f'{mode}_episodes', len(cache))
+  if mode == 'eval' or config.expl_gifs:
+    logger.video(f'{mode}_policy', video[None])
+  logger.write()
+
+
+def main(logdir, config):
+
+  logdir = os.path.join(
+      logdir, config.task, 'Ours', str(config.seed))
+
+  logdir = pathlib.Path(logdir).expanduser()
+  config.traindir = config.traindir or logdir / 'train_eps'
+  config.evaldir = config.evaldir or logdir / 'eval_eps'
+  config.steps //= config.action_repeat
+  config.eval_every //= config.action_repeat
+  config.log_every //= config.action_repeat
+  config.time_limit //= config.action_repeat
+  config.act = getattr(tf.nn, config.act)
+
+  if config.debug:
+    tf.config.experimental_run_functions_eagerly(True)
+  if config.gpu_growth:
+    message = 'No GPU found. To actually train on CPU remove this assert.'
+    assert tf.config.experimental.list_physical_devices('GPU'), message
+    for gpu in tf.config.experimental.list_physical_devices('GPU'):
+      tf.config.experimental.set_memory_growth(gpu, True)
+  assert config.precision in (16, 32), config.precision
+  if config.precision == 16:
+    prec.set_policy(prec.Policy('mixed_float16'))
+  print('Logdir', logdir)
+  logdir.mkdir(parents=True, exist_ok=True)
+  config.traindir.mkdir(parents=True, exist_ok=True)
+  config.evaldir.mkdir(parents=True, exist_ok=True)
+  step = count_steps(config.traindir)
+  logger = tools.Logger(logdir, config.action_repeat * step)
+
+  print('Create envs.')
+  if config.offline_traindir:
+    directory = config.offline_traindir.format(**vars(config))
+  else:
+    directory = config.traindir
+  train_eps = tools.load_episodes(directory, limit=config.dataset_size)
+  if config.offline_evaldir:
+    directory = config.offline_evaldir.format(**vars(config))
+  else:
+    directory = config.evaldir
+  eval_eps = tools.load_episodes(directory, limit=1)
+  make = lambda mode: make_env(config, logger, mode, train_eps, eval_eps)
+  train_envs = [make('train') for _ in range(config.envs)]
+  eval_envs = [make('eval') for _ in range(config.envs)]
+  acts = train_envs[0].action_space
+  config.num_actions = acts.n if hasattr(acts, 'n') else acts.shape[0]
+
+  prefill = max(0, config.prefill - count_steps(config.traindir))
+  print(f'Prefill dataset ({prefill} steps).')
+  random_agent = lambda o, d, s: ([acts.sample() for _ in d], s)
+  tools.simulate(random_agent, train_envs, prefill)
+  tools.simulate(random_agent, eval_envs, episodes=1)
+  logger.step = config.action_repeat * count_steps(config.traindir)
+
+  print('Simulate agent.')
+  train_dataset = make_dataset(train_eps, config)
+  eval_dataset = iter(make_dataset(eval_eps, config))
+  agent = Dreamer(config, logger, train_dataset)
+  if (logdir / 'variables.pkl').exists():
+    agent.load(logdir / 'variables.pkl')
+    agent._should_pretrain._once = False
+
+  state = None
+  suite, task = config.task.split('_', 1)
+  num_eval_episodes = 10 if suite == 'procgen' else 1
+  while agent._step.numpy().item() < config.steps:
+    logger.write()
+    print('Start evaluation.')
+    openl_joint, openl_main, openl_disen, openl_mask = agent._wm.video_pred(next(eval_dataset))
+    logger.video('eval_openl_joint', openl_joint)
+    logger.video('eval_openl_main', openl_main)
+    logger.video('eval_openl_disen', openl_disen)
+    logger.video('eval_openl_mask', openl_mask)
+    eval_policy = functools.partial(agent, training=False)
+    tools.simulate(eval_policy, eval_envs, episodes=num_eval_episodes)
+    print('Start training.')
+    state = tools.simulate(agent, train_envs, config.eval_every, state=state)
+    agent.save(logdir / 'variables.pkl')
+  for env in train_envs + eval_envs:
+    try:
+      env.close()
+    except Exception:
+      pass
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--configs', nargs='+', required=True)
+  args, remaining = parser.parse_known_args()
+  configs = yaml.safe_load((pathlib.Path(__file__).parent / 'configs.yaml').read_text())
+  config_ = {}
+  for name in args.configs:
+    config_.update(configs[name])
+  parser = argparse.ArgumentParser()
+  for key, value in config_.items():
+    arg_type = tools.args_type(value)
+    parser.add_argument(f'--{key}', type=arg_type, default=arg_type(value))
+  main(config_['logdir'], parser.parse_args(remaining))

DreamerV2/exploration.py

@@ -0,0 +1,83 @@
+import tensorflow as tf
+from tensorflow.keras.mixed_precision import experimental as prec
+from tensorflow_probability import distributions as tfd
+
+import models
+import networks
+import tools
+
+class Random(tools.Module):
+
+  def __init__(self, config):
+    self._config = config
+    self._float = prec.global_policy().comput_dtype
+
+  def actor(self, feat):
+    shape = feat.shape[:-1] + [self._config.num_actions]
+    if self._config.actor_dist == 'onehot':
+      return tools.OneHotDist(tf.zeros(shape))
+    else:
+      ones = tf.ones(shape, self._float)
+      return tfd.Uniform(-ones, ones)
+
+  def train(self, start, feat, embed, kl):
+    return None, {}
+
+
+class Plan2Explore(tools.Module):
+
+  def __init__(self, config, world_model, reward=None):
+    self._config = config
+    self._reward = reward
+    self._behavior = models.ImagBehavior(config, world_model)
+    self.actor = self._behavior.actor
+    size = {
+        'embed': 32 * config.cnn_depth,
+        'stoch': config.dyn_stoch,
+        'deter': config.dyn_deter,
+        'feat': config.dyn_stoch + config.dyn_deter,
+    }[self._config.disag_target]
+    kw = dict(
+        shape=size, layers=config.disag_layers, units=config.disag_units,
+        act=config.act)
+    self._networks = [
+        networks.DenseHead(**kw) for _ in range(config.disag_models)]
+    self._opt = tools.Optimizer(
+        'ensemble', config.model_lr, config.opt_eps, config.grad_clip,
+        config.weight_decay, opt=config.opt)
+
+  def train(self, start, feat, embed, kl):
+    metrics = {}
+    target = {
+        'embed': embed,
+        'stoch': start['stoch'],
+        'deter': start['deter'],
+        'feat': feat,
+    }[self._config.disag_target]
+    metrics.update(self._train_ensemble(feat, target))
+    metrics.update(self._behavior.train(start, self._intrinsic_reward)[-1])
+    return None, metrics
+
+  def _intrinsic_reward(self, feat, state, action):
+    preds = [head(feat, tf.float32).mean() for head in self._networks]
+    disag = tf.reduce_mean(tf.math.reduce_std(preds, 0), -1)
+    if self._config.disag_log:
+      disag = tf.math.log(disag)
+    reward = self._config.expl_intr_scale * disag
+    if self._config.expl_extr_scale:
+      reward += tf.cast(self._config.expl_extr_scale * self._reward(
+          feat, state, action), tf.float32)
+    return reward
+
+  def _train_ensemble(self, inputs, targets):
+    if self._config.disag_offset:
+      targets = targets[:, self._config.disag_offset:]
+      inputs = inputs[:, :-self._config.disag_offset]
+    targets = tf.stop_gradient(targets)
+    inputs = tf.stop_gradient(inputs)
+    with tf.GradientTape() as tape:
+      preds = [head(inputs) for head in self._networks]
+      likes = [tf.reduce_mean(pred.log_prob(targets)) for pred in preds]
+      loss = -tf.cast(tf.reduce_sum(likes), tf.float32)
+    metrics = self._opt(tape, loss, self._networks)
+    return metrics

DreamerV2/models.py

@@ -0,0 +1,429 @@
+import tensorflow as tf
+from tensorflow.keras.mixed_precision import experimental as prec
+
+import networks
+import tools
+
+
+class WorldModel(tools.Module):
+
+    def __init__(self, step, config):
+        self._step = step
+        self._config = config
+        channels = (1 if config.atari_grayscale else 3)
+        shape = config.size + (channels,)
+
+        ########
+        # Main #
+        ########
+        self.encoder = networks.ConvEncoder(
+            config.cnn_depth, config.act, config.encoder_kernels)
+        self.dynamics = networks.RSSM(
+            config.dyn_stoch, config.dyn_deter, config.dyn_hidden,
+            config.dyn_input_layers, config.dyn_output_layers, config.dyn_shared,
+            config.dyn_discrete, config.act, config.dyn_mean_act,
+            config.dyn_std_act, config.dyn_min_std, config.dyn_cell)
+        self.heads = {}
+        self.heads['reward'] = networks.DenseHead(
+            [], config.reward_layers, config.units, config.act)
+        if config.pred_discount:
+            self.heads['discount'] = networks.DenseHead(
+                [], config.discount_layers, config.units, config.act, dist='binary')
+        self._model_opt = tools.Optimizer(
+            'model', config.model_lr, config.opt_eps, config.grad_clip,
+            config.weight_decay, opt=config.opt)
+        self._scales = dict(
+            reward=config.reward_scale, discount=config.discount_scale)
+
+        #########
+        # Disen #
+        #########
+        self.disen_encoder = networks.ConvEncoder(
+            config.disen_cnn_depth, config.act, config.encoder_kernels)
+        self.disen_dynamics = networks.RSSM(
+            config.dyn_stoch, config.dyn_deter, config.dyn_hidden,
+            config.dyn_input_layers, config.dyn_output_layers, config.dyn_shared,
+            config.dyn_discrete, config.act, config.dyn_mean_act,
+            config.dyn_std_act, config.dyn_min_std, config.dyn_cell)
+
+        self.disen_heads = {}
+        self.disen_heads['reward'] = networks.DenseHead(
+            [], config.reward_layers, config.units, config.act)
+        if config.pred_discount:
+            self.disen_heads['discount'] = networks.DenseHead(
+                [], config.discount_layers, config.units, config.act, dist='binary')
+
+        self._disen_model_opt = tools.Optimizer(
+            'disen', config.model_lr, config.opt_eps, config.grad_clip,
+            config.weight_decay, opt=config.opt)
+
+        self._disen_heads_opt = {}
+        self._disen_heads_opt['reward'] = tools.Optimizer(
+            'disen_reward', config.model_lr, config.opt_eps, config.grad_clip,
+            config.weight_decay, opt=config.opt)
+        if config.pred_discount:
+            self._disen_heads_opt['discount'] = tools.Optimizer(
+                'disen_pcont', config.model_lr, config.opt_eps, config.grad_clip,
+                config.weight_decay, opt=config.opt)
+
+        # negative signs for reward/discount here
+        self._disen_scales = dict(disen_only=config.disen_only_scale,
+                                  reward=-config.disen_reward_scale, discount=-config.disen_discount_scale)
+
+        self.disen_only_image_head = networks.ConvDecoder(
+            config.disen_cnn_depth, config.act, shape, config.decoder_kernels,
+            config.decoder_thin)
+
+        ################
+        # Joint Decode #
+        ################
+        self.image_head = networks.ConvDecoderMask(
+            config.cnn_depth, config.act, shape, config.decoder_kernels,
+            config.decoder_thin)
+        self.disen_image_head = networks.ConvDecoderMask(
+            config.disen_cnn_depth, config.act, shape, config.decoder_kernels,
+            config.decoder_thin)
+        self.joint_image_head = networks.ConvDecoderMaskEnsemble(
+            self.image_head, self.disen_image_head
+        )
+
+    def train(self, data):
+        data = self.preprocess(data)
+        with tf.GradientTape() as model_tape, tf.GradientTape() as disen_tape:
+
+            # kl schedule
+            kl_balance = tools.schedule(self._config.kl_balance, self._step)
+            kl_free = tools.schedule(self._config.kl_free, self._step)
+            kl_scale = tools.schedule(self._config.kl_scale, self._step)
+
+            # Main
+            embed = self.encoder(data)
+            post, prior = self.dynamics.observe(embed, data['action'])
+            kl_loss, kl_value = self.dynamics.kl_loss(
+                post, prior, kl_balance, kl_free, kl_scale)
+            feat = self.dynamics.get_feat(post)
+            likes = {}
+            for name, head in self.heads.items():
+                grad_head = (name in self._config.grad_heads)
+                inp = feat if grad_head else tf.stop_gradient(feat)
+                pred = head(inp, tf.float32)
+                like = pred.log_prob(tf.cast(data[name], tf.float32))
+                likes[name] = tf.reduce_mean(
+                    like) * self._scales.get(name, 1.0)
+
+            # Disen
+            embed_disen = self.disen_encoder(data)
+            post_disen, prior_disen = self.disen_dynamics.observe(
+                embed_disen, data['action'])
+            kl_loss_disen, kl_value_disen = self.dynamics.kl_loss(
+                post_disen, prior_disen, kl_balance, kl_free, kl_scale)
+            feat_disen = self.disen_dynamics.get_feat(post_disen)
+
+            # Optimize disen reward/pcont till optimal
+            disen_metrics = dict(reward={}, discount={})
+            loss_disen = dict(reward=None, discount=None)
+            for _ in range(self._config.num_reward_opt_iters):
+                with tf.GradientTape() as disen_reward_tape, tf.GradientTape() as disen_pcont_tape:
+                    disen_gradient_tapes = dict(
+                        reward=disen_reward_tape, discount=disen_pcont_tape)
+                    for name, head in self.disen_heads.items():
+                        pred_disen = head(
+                            tf.stop_gradient(feat_disen), tf.float32)
+                        loss_disen[name] = -tf.reduce_mean(pred_disen.log_prob(
+                            tf.cast(data[name], tf.float32)))
+                for name, head in self.disen_heads.items():
+                    disen_metrics[name] = self._disen_heads_opt[name](
+                        disen_gradient_tapes[name], loss_disen[name], [head], prefix='disen_neg')
+
+            # Compute likes for disen model (including negative gradients)
+            likes_disen = {}
+            for name, head in self.disen_heads.items():
+                pred_disen = head(feat_disen, tf.float32)
+                like_disen = pred_disen.log_prob(
+                    tf.cast(data[name], tf.float32))
+                likes_disen[name] = tf.reduce_mean(
+                    like_disen) * self._disen_scales.get(name, -1.0)
+            disen_only_image_pred = self.disen_only_image_head(
+                feat_disen, tf.float32)
+            disen_only_image_like = tf.reduce_mean(disen_only_image_pred.log_prob(
+                tf.cast(data['image'], tf.float32))) * self._disen_scales.get('disen_only', 1.0)
+            likes_disen['disen_only'] = disen_only_image_like
+
+            # Joint decode
+            image_pred_joint, _, _, _ = self.joint_image_head(
+                feat, feat_disen, tf.float32)
+            image_like = tf.reduce_mean(image_pred_joint.log_prob(
+                tf.cast(data['image'], tf.float32)))
+            likes['image'] = image_like
+            likes_disen['image'] = image_like
+
+            # Compute loss
+            model_loss = kl_loss - sum(likes.values())
+            disen_loss = kl_loss_disen - sum(likes_disen.values())
+
+        model_parts = [self.encoder, self.dynamics,
+                       self.joint_image_head] + list(self.heads.values())
+        disen_parts = [self.disen_encoder, self.disen_dynamics,
+                       self.joint_image_head, self.disen_only_image_head]
+
+        metrics = self._model_opt(
+            model_tape, model_loss, model_parts, prefix='main')
+        disen_model_metrics = self._disen_model_opt(
+            disen_tape, disen_loss, disen_parts, prefix='disen')
+
+        metrics['kl_balance'] = kl_balance
+        metrics['kl_free'] = kl_free
+        metrics['kl_scale'] = kl_scale
+        metrics.update({f'{name}_loss': -like for name,
+                        like in likes.items()})
+
+        metrics['disen/disen_only_image_loss'] = -disen_only_image_like
+        metrics['disen/disen_reward_loss'] = -likes_disen['reward'] / \
+            self._disen_scales.get('reward', -1.0)
+        metrics['disen/disen_discount_loss'] = -likes_disen['discount'] / \
+            self._disen_scales.get('discount', -1.0)
+
+        metrics['kl'] = tf.reduce_mean(kl_value)
+        metrics['prior_ent'] = self.dynamics.get_dist(prior).entropy()
+        metrics['post_ent'] = self.dynamics.get_dist(post).entropy()
+        metrics['disen/kl'] = tf.reduce_mean(kl_value_disen)
+        metrics['disen/prior_ent'] = self.dynamics.get_dist(
+            prior_disen).entropy()
+        metrics['disen/post_ent'] = self.dynamics.get_dist(
+            post_disen).entropy()
+
+        metrics.update(
+            {f'{key}': value for key, value in disen_metrics['reward'].items()})
+        metrics.update(
+            {f'{key}': value for key, value in disen_metrics['discount'].items()})
+        metrics.update(
+            {f'{key}': value for key, value in disen_model_metrics.items()})
+
+        return embed, post, feat, kl_value, metrics
+
+    @tf.function
+    def preprocess(self, obs):
+        dtype = prec.global_policy().compute_dtype
+        obs = obs.copy()
+        obs['image'] = tf.cast(obs['image'], dtype) / 255.0 - 0.5
+        obs['reward'] = getattr(tf, self._config.clip_rewards)(obs['reward'])
+        if 'discount' in obs:
+            obs['discount'] *= self._config.discount
+        for key, value in obs.items():
+            if tf.dtypes.as_dtype(value.dtype) in (
+                    tf.float16, tf.float32, tf.float64):
+                obs[key] = tf.cast(value, dtype)
+        return obs
+
+    @tf.function
+    def video_pred(self, data):
+        data = self.preprocess(data)
+        truth = data['image'][:6] + 0.5
+
+        embed = self.encoder(data)
+        embed_disen = self.disen_encoder(data)
+        states, _ = self.dynamics.observe(
+            embed[:6, :5], data['action'][:6, :5])
+        states_disen, _ = self.disen_dynamics.observe(
+            embed_disen[:6, :5], data['action'][:6, :5])
+        feats = self.dynamics.get_feat(states)
+        feats_disen = self.disen_dynamics.get_feat(states_disen)
+        recon_joint, recon_main, recon_disen, recon_mask = self.joint_image_head(
+            feats, feats_disen)
+        recon_joint = recon_joint.mode()[:6]
+        recon_main = recon_main.mode()[:6]
+        recon_disen = recon_disen.mode()[:6]
+        recon_mask = recon_mask[:6]
+
+        init = {k: v[:, -1] for k, v in states.items()}
+        init_disen = {k: v[:, -1] for k, v in states_disen.items()}
+        prior = self.dynamics.imagine(
+            data['action'][:6, 5:], init)
+        prior_disen = self.disen_dynamics.imagine(
+            data['action'][:6, 5:], init_disen)
+        _feats = self.dynamics.get_feat(prior)
+        _feats_disen = self.disen_dynamics.get_feat(prior_disen)
+        openl_joint, openl_main, openl_disen, openl_mask = self.joint_image_head(
+            _feats, _feats_disen)
+        openl_joint = openl_joint.mode()
+        openl_main = openl_main.mode()
+        openl_disen = openl_disen.mode()
+
+        model_joint = tf.concat(
+            [recon_joint[:, :5] + 0.5, openl_joint + 0.5], 1)
+        error_joint = (model_joint - truth + 1) / 2
+        model_main = tf.concat(
+            [recon_main[:, :5] + 0.5, openl_main + 0.5], 1)
+        error_main = (model_main - truth + 1) / 2
+        model_disen = tf.concat(
+            [recon_disen[:, :5] + 0.5, openl_disen + 0.5], 1)
+        error_disen = (model_disen - truth + 1) / 2
+        model_mask = tf.concat(
+            [recon_mask[:, :5] + 0.5, openl_mask + 0.5], 1)
+
+        output_joint = tf.concat([truth, model_joint, error_joint], 2)
+        output_main = tf.concat([truth, model_main, error_main], 2)
+        output_disen = tf.concat([truth, model_disen, error_disen], 2)
+        output_mask = model_mask
+
+        return output_joint, output_main, output_disen, output_mask
+
+
+class ImagBehavior(tools.Module):
+
+    def __init__(self, config, world_model, stop_grad_actor=True, reward=None):
+        self._config = config
+        self._world_model = world_model
+        self._stop_grad_actor = stop_grad_actor
+        self._reward = reward
+        self.actor = networks.ActionHead(
+            config.num_actions, config.actor_layers, config.units, config.act,
+            config.actor_dist, config.actor_init_std, config.actor_min_std,
+            config.actor_dist, config.actor_temp, config.actor_outscale)
+        self.value = networks.DenseHead(
+            [], config.value_layers, config.units, config.act,
+            config.value_head)
+        if config.slow_value_target or config.slow_actor_target:
+            self._slow_value = networks.DenseHead(
+                [], config.value_layers, config.units, config.act)
+            self._updates = tf.Variable(0, tf.int64)
+        kw = dict(wd=config.weight_decay, opt=config.opt)
+        self._actor_opt = tools.Optimizer(
+            'actor', config.actor_lr, config.opt_eps, config.actor_grad_clip, **kw)
+        self._value_opt = tools.Optimizer(
+            'value', config.value_lr, config.opt_eps, config.value_grad_clip, **kw)
+
+    def train(
+            self, start, objective=None, imagine=None, tape=None, repeats=None):
+        objective = objective or self._reward
+        self._update_slow_target()
+        metrics = {}
+        with (tape or tf.GradientTape()) as actor_tape:
+            assert bool(objective) != bool(imagine)
+            if objective:
+                imag_feat, imag_state, imag_action = self._imagine(
+                    start, self.actor, self._config.imag_horizon, repeats)
+                reward = objective(imag_feat, imag_state, imag_action)
+            else:
+                imag_feat, imag_state, imag_action, reward = imagine(start)
+            actor_ent = self.actor(imag_feat, tf.float32).entropy()
+            state_ent = self._world_model.dynamics.get_dist(
+                imag_state, tf.float32).entropy()
+            target, weights = self._compute_target(
+                imag_feat, reward, actor_ent, state_ent,
+                self._config.slow_actor_target)
+            actor_loss, mets = self._compute_actor_loss(
+                imag_feat, imag_state, imag_action, target, actor_ent, state_ent,
+                weights)
+            metrics.update(mets)
+        if self._config.slow_value_target != self._config.slow_actor_target:
+            target, weights = self._compute_target(
+                imag_feat, reward, actor_ent, state_ent,
+                self._config.slow_value_target)
+        with tf.GradientTape() as value_tape:
+            value = self.value(imag_feat, tf.float32)[:-1]
+            value_loss = -value.log_prob(tf.stop_gradient(target))
+            if self._config.value_decay:
+                value_loss += self._config.value_decay * value.mode()
+            value_loss = tf.reduce_mean(weights[:-1] * value_loss)
+        metrics['reward_mean'] = tf.reduce_mean(reward)
+        metrics['reward_std'] = tf.math.reduce_std(reward)
+        metrics['actor_ent'] = tf.reduce_mean(actor_ent)
+        metrics.update(self._actor_opt(actor_tape, actor_loss, [self.actor]))
+        metrics.update(self._value_opt(value_tape, value_loss, [self.value]))
+        return imag_feat, imag_state, imag_action, weights, metrics
+
+    def _imagine(self, start, policy, horizon, repeats=None):
+        dynamics = self._world_model.dynamics
+        if repeats:
+            start = {k: tf.repeat(v, repeats, axis=1)
+                     for k, v in start.items()}
+
+        def flatten(x): return tf.reshape(x, [-1] + list(x.shape[2:]))
+        start = {k: flatten(v) for k, v in start.items()}
+
+        def step(prev, _):
+            state, _, _ = prev
+            feat = dynamics.get_feat(state)
+            inp = tf.stop_gradient(feat) if self._stop_grad_actor else feat
+            action = policy(inp).sample()
+            succ = dynamics.img_step(
+                state, action, sample=self._config.imag_sample)
+            return succ, feat, action
+        feat = 0 * dynamics.get_feat(start)
+        action = policy(feat).mode()
+        succ, feats, actions = tools.static_scan(
+            step, tf.range(horizon), (start, feat, action))
+        states = {k: tf.concat([
+            start[k][None], v[:-1]], 0) for k, v in succ.items()}
+        if repeats:
+            def unfold(tensor):
+                s = tensor.shape
+                return tf.reshape(tensor, [s[0], s[1] // repeats, repeats] + s[2:])
+            states, feats, actions = tf.nest.map_structure(
+                unfold, (states, feats, actions))
+        return feats, states, actions
+
+    def _compute_target(self, imag_feat, reward, actor_ent, state_ent, slow):
+        reward = tf.cast(reward, tf.float32)
+        if 'discount' in self._world_model.heads:
+            discount = self._world_model.heads['discount'](
+                imag_feat, tf.float32).mean()
+        else:
+            discount = self._config.discount * tf.ones_like(reward)
+        if self._config.future_entropy and tf.greater(
+                self._config.actor_entropy(), 0):
+            reward += self._config.actor_entropy() * actor_ent
+        if self._config.future_entropy and tf.greater(
+                self._config.actor_state_entropy(), 0):
+            reward += self._config.actor_state_entropy() * state_ent
+        if slow:
+            value = self._slow_value(imag_feat, tf.float32).mode()
+        else:
+            value = self.value(imag_feat, tf.float32).mode()
+        target = tools.lambda_return(
+            reward[:-1], value[:-1], discount[:-1],
+            bootstrap=value[-1], lambda_=self._config.discount_lambda, axis=0)
+        weights = tf.stop_gradient(tf.math.cumprod(tf.concat(
+            [tf.ones_like(discount[:1]), discount[:-1]], 0), 0))
+        return target, weights
+
+    def _compute_actor_loss(
+            self, imag_feat, imag_state, imag_action, target, actor_ent, state_ent,
+            weights):
+        metrics = {}
+        inp = tf.stop_gradient(
+            imag_feat) if self._stop_grad_actor else imag_feat
+        policy = self.actor(inp, tf.float32)
+        actor_ent = policy.entropy()
+        if self._config.imag_gradient == 'dynamics':
+            actor_target = target
+        elif self._config.imag_gradient == 'reinforce':
+            imag_action = tf.cast(imag_action, tf.float32)
+            actor_target = policy.log_prob(imag_action)[:-1] * tf.stop_gradient(
+                target - self.value(imag_feat[:-1], tf.float32).mode())
+        elif self._config.imag_gradient == 'both':
+            imag_action = tf.cast(imag_action, tf.float32)
+            actor_target = policy.log_prob(imag_action)[:-1] * tf.stop_gradient(
+                target - self.value(imag_feat[:-1], tf.float32).mode())
+            mix = self._config.imag_gradient_mix()
+            actor_target = mix * target + (1 - mix) * actor_target
+            metrics['imag_gradient_mix'] = mix
+        else:
+            raise NotImplementedError(self._config.imag_gradient)
+        if not self._config.future_entropy and tf.greater(
+                self._config.actor_entropy(), 0):
+            actor_target += self._config.actor_entropy() * actor_ent[:-1]
+        if not self._config.future_entropy and tf.greater(
+                self._config.actor_state_entropy(), 0):
+            actor_target += self._config.actor_state_entropy() * state_ent[:-1]
+        actor_loss = -tf.reduce_mean(weights[:-1] * actor_target)
+        return actor_loss, metrics
+
+    def _update_slow_target(self):
+        if self._config.slow_value_target or self._config.slow_actor_target:
+            if self._updates % self._config.slow_target_update == 0:
+                mix = self._config.slow_target_fraction
+                for s, d in zip(self.value.variables, self._slow_value.variables):
+                    d.assign(mix * s + (1 - mix) * d)
+            self._updates.assign_add(1)

DreamerV2/networks.py

@@ -0,0 +1,465 @@
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras import layers as tfkl
+from tensorflow_probability import distributions as tfd
+from tensorflow.keras.mixed_precision import experimental as prec
+
+import tools
+
+class RSSM(tools.Module):
+
+    def __init__(
+            self, stoch=30, deter=200, hidden=200, layers_input=1, layers_output=1,
+            shared=False, discrete=False, act=tf.nn.elu, mean_act='none',
+            std_act='softplus', min_std=0.1, cell='keras'):
+        super().__init__()
+        self._stoch = stoch
+        self._deter = deter
+        self._hidden = hidden
+        self._min_std = min_std
+        self._layers_input = layers_input
+        self._layers_output = layers_output
+        self._shared = shared
+        self._discrete = discrete
+        self._act = act
+        self._mean_act = mean_act
+        self._std_act = std_act
+        self._embed = None
+        if cell == 'gru':
+            self._cell = tfkl.GRUCell(self._deter)
+        elif cell == 'gru_layer_norm':
+            self._cell = GRUCell(self._deter, norm=True)
+        else:
+            raise NotImplementedError(cell)
+
+    def initial(self, batch_size):
+        dtype = prec.global_policy().compute_dtype
+        if self._discrete:
+            state = dict(
+                logit=tf.zeros(
+                    [batch_size, self._stoch, self._discrete], dtype),
+                stoch=tf.zeros(
+                    [batch_size, self._stoch, self._discrete], dtype),
+                deter=self._cell.get_initial_state(None, batch_size, dtype))
+        else:
+            state = dict(
+                mean=tf.zeros([batch_size, self._stoch], dtype),
+                std=tf.zeros([batch_size, self._stoch], dtype),
+                stoch=tf.zeros([batch_size, self._stoch], dtype),
+                deter=self._cell.get_initial_state(None, batch_size, dtype))
+        return state
+
+    @tf.function
+    def observe(self, embed, action, state=None):
+        def swap(x): return tf.transpose(
+            x, [1, 0] + list(range(2, len(x.shape))))
+        if state is None:
+            state = self.initial(tf.shape(action)[0])
+        embed, action = swap(embed), swap(action)
+        post, prior = tools.static_scan(
+            lambda prev, inputs: self.obs_step(prev[0], *inputs),
+            (action, embed), (state, state))
+        post = {k: swap(v) for k, v in post.items()}
+        prior = {k: swap(v) for k, v in prior.items()}
+        return post, prior
+
+    @tf.function
+    def imagine(self, action, state=None):
+        def swap(x): return tf.transpose(
+            x, [1, 0] + list(range(2, len(x.shape))))
+        if state is None:
+            state = self.initial(tf.shape(action)[0])
+        assert isinstance(state, dict), state
+        action = swap(action)
+        prior = tools.static_scan(self.img_step, action, state)
+        prior = {k: swap(v) for k, v in prior.items()}
+        return prior
+
+    def get_feat(self, state):
+        stoch = state['stoch']
+        if self._discrete:
+            shape = stoch.shape[:-2] + [self._stoch * self._discrete]
+            stoch = tf.reshape(stoch, shape)
+        return tf.concat([stoch, state['deter']], -1)
+
+    def get_dist(self, state, dtype=None):
+        if self._discrete:
+            logit = state['logit']
+            logit = tf.cast(logit, tf.float32)
+            dist = tfd.Independent(tools.OneHotDist(logit), 1)
+            if dtype != tf.float32:
+                dist = tools.DtypeDist(dist, dtype or state['logit'].dtype)
+        else:
+            mean, std = state['mean'], state['std']
+            if dtype:
+                mean = tf.cast(mean, dtype)
+                std = tf.cast(std, dtype)
+            dist = tfd.MultivariateNormalDiag(mean, std)
+        return dist
+
+    @tf.function
+    def obs_step(self, prev_state, prev_action, embed, sample=True):
+        if not self._embed:
+            self._embed = embed.shape[-1]
+        prior = self.img_step(prev_state, prev_action, None, sample)
+        if self._shared:
+            post = self.img_step(prev_state, prev_action, embed, sample)
+        else:
+            x = tf.concat([prior['deter'], embed], -1)
+            for i in range(self._layers_output):
+                x = self.get(f'obi{i}', tfkl.Dense, self._hidden, self._act)(x)
+            stats = self._suff_stats_layer('obs', x)
+            if sample:
+                stoch = self.get_dist(stats).sample()
+            else:
+                stoch = self.get_dist(stats).mode()
+            post = {'stoch': stoch, 'deter': prior['deter'], **stats}
+        return post, prior
+
+    @tf.function
+    def img_step(self, prev_state, prev_action, embed=None, sample=True):
+        prev_stoch = prev_state['stoch']
+        if self._discrete:
+            shape = prev_stoch.shape[:-2] + [self._stoch * self._discrete]
+            prev_stoch = tf.reshape(prev_stoch, shape)
+        if self._shared:
+            if embed is None:
+                shape = prev_action.shape[:-1] + [self._embed]
+                embed = tf.zeros(shape, prev_action.dtype)
+            x = tf.concat([prev_stoch, prev_action, embed], -1)
+        else:
+            x = tf.concat([prev_stoch, prev_action], -1)
+        for i in range(self._layers_input):
+            x = self.get(f'ini{i}', tfkl.Dense, self._hidden, self._act)(x)
+        x, deter = self._cell(x, [prev_state['deter']])
+        deter = deter[0]  # Keras wraps the state in a list.
+        for i in range(self._layers_output):
+            x = self.get(f'imo{i}', tfkl.Dense, self._hidden, self._act)(x)
+        stats = self._suff_stats_layer('ims', x)
+        if sample:
+            stoch = self.get_dist(stats).sample()
+        else:
+            stoch = self.get_dist(stats).mode()
+        prior = {'stoch': stoch, 'deter': deter, **stats}
+        return prior
+
+    def _suff_stats_layer(self, name, x):
+        if self._discrete:
+            x = self.get(name, tfkl.Dense, self._stoch *
+                         self._discrete, None)(x)
+            logit = tf.reshape(x, x.shape[:-1] + [self._stoch, self._discrete])
+            return {'logit': logit}
+        else:
+            x = self.get(name, tfkl.Dense, 2 * self._stoch, None)(x)
+            mean, std = tf.split(x, 2, -1)
+            mean = {
+                'none': lambda: mean,
+                'tanh5': lambda: 5.0 * tf.math.tanh(mean / 5.0),
+            }[self._mean_act]()
+            std = {
+                'softplus': lambda: tf.nn.softplus(std),
+                'abs': lambda: tf.math.abs(std + 1),
+                'sigmoid': lambda: tf.nn.sigmoid(std),
+                'sigmoid2': lambda: 2 * tf.nn.sigmoid(std / 2),
+            }[self._std_act]()
+            std = std + self._min_std
+            return {'mean': mean, 'std': std}
+
+    def kl_loss(self, post, prior, balance, free, scale):
+        kld = tfd.kl_divergence
+        def dist(x): return self.get_dist(x, tf.float32)
+        if balance == 0.5:
+            value = kld(dist(prior), dist(post))
+            loss = tf.reduce_mean(tf.maximum(value, free))
+        else:
+            def sg(x): return tf.nest.map_structure(tf.stop_gradient, x)
+            value = kld(dist(prior), dist(sg(post)))
+            pri = tf.reduce_mean(value)
+            pos = tf.reduce_mean(kld(dist(sg(prior)), dist(post)))
+            pri, pos = tf.maximum(pri, free), tf.maximum(pos, free)
+            loss = balance * pri + (1 - balance) * pos
+        loss *= scale
+        return loss, value
+
+
+class ConvEncoder(tools.Module):
+
+    def __init__(
+            self, depth=32, act=tf.nn.relu, kernels=(4, 4, 4, 4)):
+        self._act = act
+        self._depth = depth
+        self._kernels = kernels
+
+    def __call__(self, obs):
+        kwargs = dict(strides=2, activation=self._act)
+        Conv = tfkl.Conv2D
+        x = tf.reshape(obs['image'], (-1,) + tuple(obs['image'].shape[-3:]))
+        x = self.get('h1', Conv, 1 * self._depth,
+                     self._kernels[0], **kwargs)(x)
+        x = self.get('h2', Conv, 2 * self._depth,
+                     self._kernels[1], **kwargs)(x)
+        x = self.get('h3', Conv, 4 * self._depth,
+                     self._kernels[2], **kwargs)(x)
+        x = self.get('h4', Conv, 8 * self._depth,
+                     self._kernels[3], **kwargs)(x)
+        x = tf.reshape(x, [x.shape[0], np.prod(x.shape[1:])])
+        shape = tf.concat([tf.shape(obs['image'])[:-3], [x.shape[-1]]], 0)
+        return tf.reshape(x, shape)
+
+
+class ConvDecoder(tools.Module):
+
+    def __init__(
+            self, depth=32, act=tf.nn.relu, shape=(64, 64, 3), kernels=(5, 5, 6, 6),
+            thin=True):
+        self._act = act
+        self._depth = depth
+        self._shape = shape
+        self._kernels = kernels
+        self._thin = thin
+
+    def __call__(self, features, dtype=None):
+        kwargs = dict(strides=2, activation=self._act)
+        ConvT = tfkl.Conv2DTranspose
+        if self._thin:
+            x = self.get('h1', tfkl.Dense, 32 * self._depth, None)(features)
+            x = tf.reshape(x, [-1, 1, 1, 32 * self._depth])
+        else:
+            x = self.get('h1', tfkl.Dense, 128 * self._depth, None)(features)
+            x = tf.reshape(x, [-1, 2, 2, 32 * self._depth])
+        x = self.get('h2', ConvT, 4 * self._depth,
+                     self._kernels[0], **kwargs)(x)
+        x = self.get('h3', ConvT, 2 * self._depth,
+                     self._kernels[1], **kwargs)(x)
+        x = self.get('h4', ConvT, 1 * self._depth,
+                     self._kernels[2], **kwargs)(x)
+        x = self.get(
+            'h5', ConvT, self._shape[-1], self._kernels[3], strides=2)(x)
+        mean = tf.reshape(x, tf.concat(
+            [tf.shape(features)[:-1], self._shape], 0))
+        if dtype:
+            mean = tf.cast(mean, dtype)
+        return tfd.Independent(tfd.Normal(mean, 1), len(self._shape))
+
+
+class ConvDecoderMask(tools.Module):
+
+    def __init__(
+            self, depth=32, act=tf.nn.relu, shape=(64, 64, 3), kernels=(5, 5, 6, 6),
+            thin=True):
+        self._act = act
+        self._depth = depth
+        self._shape = shape
+        self._kernels = kernels
+        self._thin = thin
+
+    def __call__(self, features, dtype=None):
+        kwargs = dict(strides=2, activation=self._act)
+        ConvT = tfkl.Conv2DTranspose
+        if self._thin:
+            x = self.get('h1', tfkl.Dense, 32 * self._depth, None)(features)
+            x = tf.reshape(x, [-1, 1, 1, 32 * self._depth])
+        else:
+            x = self.get('h1', tfkl.Dense, 128 * self._depth, None)(features)
+            x = tf.reshape(x, [-1, 2, 2, 32 * self._depth])
+        x = self.get('h2', ConvT, 4 * self._depth,
+                     self._kernels[0], **kwargs)(x)
+        x = self.get('h3', ConvT, 2 * self._depth,
+                     self._kernels[1], **kwargs)(x)
+        x = self.get('h4', ConvT, 1 * self._depth,
+                     self._kernels[2], **kwargs)(x)
+        x = self.get(
+            'h5', ConvT, 2 * self._shape[-1], self._kernels[3], strides=2)(x)
+        mean, mask = tf.split(x, [self._shape[-1], self._shape[-1]], -1)
+        mean = tf.reshape(mean, tf.concat(
+            [tf.shape(features)[:-1], self._shape], 0))
+        mask = tf.reshape(mask, tf.concat(
+            [tf.shape(features)[:-1], self._shape], 0))
+        if dtype:
+            mean = tf.cast(mean, dtype)
+            mask = tf.cast(mask, dtype)
+        return tfd.Independent(tfd.Normal(mean, 1), len(self._shape)), mask
+      
+
+class ConvDecoderMaskEnsemble(tools.Module):
+    """
+    ensemble two convdecoder with <Normal, mask> outputs
+    NOTE: remove pred1/pred2 for maximum performance.
+    """
+
+    def __init__(self, decoder1, decoder2):
+        self._decoder1 = decoder1
+        self._decoder2 = decoder2
+        self._shape = decoder1._shape
+
+    def __call__(self, feat1, feat2, dtype=None):
+        kwargs = dict(strides=1, activation=tf.nn.sigmoid)
+        pred1, mask1 = self._decoder1(feat1, dtype)
+        pred2, mask2 = self._decoder2(feat2, dtype)
+        mean1 = pred1.submodules[0].loc
+        mean2 = pred2.submodules[0].loc
+        mask_feat = tf.concat([mask1, mask2], -1)
+        mask = self.get('mask1', tfkl.Conv2D, 1, 1, **kwargs)(mask_feat)
+        mask_use1 = mask
+        mask_use2 = 1-mask
+        mean = mean1 * tf.cast(mask_use1, mean1.dtype) + \
+            mean2 * tf.cast(mask_use2, mean2.dtype)
+        return tfd.Independent(tfd.Normal(mean, 1), len(self._shape)), pred1, pred2, tf.cast(mask_use1, mean1.dtype)
+
+
+class DenseHead(tools.Module):
+
+    def __init__(
+            self, shape, layers, units, act=tf.nn.elu, dist='normal', std=1.0):
+        self._shape = (shape,) if isinstance(shape, int) else shape
+        self._layers = layers
+        self._units = units
+        self._act = act
+        self._dist = dist
+        self._std = std
+
+    def __call__(self, features, dtype=None):
+        x = features
+        for index in range(self._layers):
+            x = self.get(f'h{index}', tfkl.Dense, self._units, self._act)(x)
+        mean = self.get(f'hmean', tfkl.Dense, np.prod(self._shape))(x)
+        mean = tf.reshape(mean, tf.concat(
+            [tf.shape(features)[:-1], self._shape], 0))
+        if self._std == 'learned':
+            std = self.get(f'hstd', tfkl.Dense, np.prod(self._shape))(x)
+            std = tf.nn.softplus(std) + 0.01
+            std = tf.reshape(std, tf.concat(
+                [tf.shape(features)[:-1], self._shape], 0))
+        else:
+            std = self._std
+        if dtype:
+            mean, std = tf.cast(mean, dtype), tf.cast(std, dtype)
+        if self._dist == 'normal':
+            return tfd.Independent(tfd.Normal(mean, std), len(self._shape))
+        if self._dist == 'huber':
+            return tfd.Independent(
+                tools.UnnormalizedHuber(mean, std, 1.0), len(self._shape))
+        if self._dist == 'binary':
+            return tfd.Independent(tfd.Bernoulli(mean), len(self._shape))
+        raise NotImplementedError(self._dist)
+
+
+class ActionHead(tools.Module):
+
+    def __init__(
+            self, size, layers, units, act=tf.nn.elu, dist='trunc_normal',
+            init_std=0.0, min_std=0.1, action_disc=5, temp=0.1, outscale=0):
+        # assert min_std <= 2
+        self._size = size
+        self._layers = layers
+        self._units = units
+        self._dist = dist
+        self._act = act
+        self._min_std = min_std
+        self._init_std = init_std
+        self._action_disc = action_disc
+        self._temp = temp() if callable(temp) else temp
+        self._outscale = outscale
+
+    def __call__(self, features, dtype=None):
+        x = features
+        for index in range(self._layers):
+            kw = {}
+            if index == self._layers - 1 and self._outscale:
+                kw['kernel_initializer'] = tf.keras.initializers.VarianceScaling(
+                    self._outscale)
+            x = self.get(f'h{index}', tfkl.Dense,
+                         self._units, self._act, **kw)(x)
+        if self._dist == 'tanh_normal':
+            # https://www.desmos.com/calculator/rcmcf5jwe7
+            x = self.get(f'hout', tfkl.Dense, 2 * self._size)(x)
+            if dtype:
+                x = tf.cast(x, dtype)
+            mean, std = tf.split(x, 2, -1)
+            mean = tf.tanh(mean)
+            std = tf.nn.softplus(std + self._init_std) + self._min_std
+            dist = tfd.Normal(mean, std)
+            dist = tfd.TransformedDistribution(dist, tools.TanhBijector())
+            dist = tfd.Independent(dist, 1)
+            dist = tools.SampleDist(dist)
+        elif self._dist == 'tanh_normal_5':
+            x = self.get(f'hout', tfkl.Dense, 2 * self._size)(x)
+            if dtype:
+                x = tf.cast(x, dtype)
+            mean, std = tf.split(x, 2, -1)
+            mean = 5 * tf.tanh(mean / 5)
+            std = tf.nn.softplus(std + 5) + 5
+            dist = tfd.Normal(mean, std)
+            dist = tfd.TransformedDistribution(dist, tools.TanhBijector())
+            dist = tfd.Independent(dist, 1)
+            dist = tools.SampleDist(dist)
+        elif self._dist == 'normal':
+            x = self.get(f'hout', tfkl.Dense, 2 * self._size)(x)
+            if dtype:
+                x = tf.cast(x, dtype)
+            mean, std = tf.split(x, 2, -1)
+            std = tf.nn.softplus(std + self._init_std) + self._min_std
+            dist = tfd.Normal(mean, std)
+            dist = tfd.Independent(dist, 1)
+        elif self._dist == 'normal_1':
+            mean = self.get(f'hout', tfkl.Dense, self._size)(x)
+            if dtype:
+                mean = tf.cast(mean, dtype)
+            dist = tfd.Normal(mean, 1)
+            dist = tfd.Independent(dist, 1)
+        elif self._dist == 'trunc_normal':
+            # https://www.desmos.com/calculator/mmuvuhnyxo
+            x = self.get(f'hout', tfkl.Dense, 2 * self._size)(x)
+            x = tf.cast(x, tf.float32)
+            mean, std = tf.split(x, 2, -1)
+            mean = tf.tanh(mean)
+            std = 2 * tf.nn.sigmoid(std / 2) + self._min_std
+            dist = tools.SafeTruncatedNormal(mean, std, -1, 1)
+            dist = tools.DtypeDist(dist, dtype)
+            dist = tfd.Independent(dist, 1)
+        elif self._dist == 'onehot':
+            x = self.get(f'hout', tfkl.Dense, self._size)(x)
+            x = tf.cast(x, tf.float32)
+            dist = tools.OneHotDist(x, dtype=dtype)
+            dist = tools.DtypeDist(dist, dtype)
+        elif self._dist == 'onehot_gumble':
+            x = self.get(f'hout', tfkl.Dense, self._size)(x)
+            if dtype:
+                x = tf.cast(x, dtype)
+            temp = self._temp
+            dist = tools.GumbleDist(temp, x, dtype=dtype)
+        else:
+            raise NotImplementedError(self._dist)
+        return dist
+
+
+class GRUCell(tf.keras.layers.AbstractRNNCell):
+
+    def __init__(self, size, norm=False, act=tf.tanh, update_bias=-1, **kwargs):
+        super().__init__()
+        self._size = size
+        self._act = act
+        self._norm = norm
+        self._update_bias = update_bias
+        self._layer = tfkl.Dense(3 * size, use_bias=norm is not None, **kwargs)
+        if norm:
+            self._norm = tfkl.LayerNormalization(dtype=tf.float32)
+
+    @property
+    def state_size(self):
+        return self._size
+
+    def call(self, inputs, state):
+        state = state[0]  # Keras wraps the state in a list.
+        parts = self._layer(tf.concat([inputs, state], -1))
+        if self._norm:
+            dtype = parts.dtype
+            parts = tf.cast(parts, tf.float32)
+            parts = self._norm(parts)
+            parts = tf.cast(parts, dtype)
+        reset, cand, update = tf.split(parts, 3, -1)
+        reset = tf.nn.sigmoid(reset)
+        cand = self._act(reset * cand)
+        update = tf.nn.sigmoid(update + self._update_bias)
+        output = update * cand + (1 - update) * state
+        return output, [output]

DreamerV2/tools.py

@@ -0,0 +1,694 @@
+import datetime
+import io
+import json
+import pathlib
+import pickle
+import re
+import time
+import uuid
+
+import numpy as np
+import tensorflow as tf
+import tensorflow.compat.v1 as tf1
+import tensorflow_probability as tfp
+from tensorflow.keras.mixed_precision import experimental as prec
+from tensorflow_probability import distributions as tfd
+
+
+# Patch to ignore seed to avoid synchronization across GPUs.
+_orig_random_categorical = tf.random.categorical
+def random_categorical(*args, **kwargs):
+  kwargs['seed'] = None
+  return _orig_random_categorical(*args, **kwargs)
+tf.random.categorical = random_categorical
+
+# Patch to ignore seed to avoid synchronization across GPUs.
+_orig_random_normal = tf.random.normal
+def random_normal(*args, **kwargs):
+  kwargs['seed'] = None
+  return _orig_random_normal(*args, **kwargs)
+tf.random.normal = random_normal
+
+
+class AttrDict(dict):
+
+  __setattr__ = dict.__setitem__
+  __getattr__ = dict.__getitem__
+
+
+class Module(tf.Module):
+
+  def save(self, filename):
+    values = tf.nest.map_structure(lambda x: x.numpy(), self.variables)
+    amount = len(tf.nest.flatten(values))
+    count = int(sum(np.prod(x.shape) for x in tf.nest.flatten(values)))
+    print(f'Save checkpoint with {amount} tensors and {count} parameters.')
+    with pathlib.Path(filename).open('wb') as f:
+      pickle.dump(values, f)
+
+  def load(self, filename):
+    with pathlib.Path(filename).open('rb') as f:
+      values = pickle.load(f)
+    amount = len(tf.nest.flatten(values))
+    count = int(sum(np.prod(x.shape) for x in tf.nest.flatten(values)))
+    print(f'Load checkpoint with {amount} tensors and {count} parameters.')
+    tf.nest.map_structure(lambda x, y: x.assign(y), self.variables, values)
+
+  def get(self, name, ctor, *args, **kwargs):
+    # Create or get layer by name to avoid mentioning it in the constructor.
+    if not hasattr(self, '_modules'):
+      self._modules = {}
+    if name not in self._modules:
+      self._modules[name] = ctor(*args, **kwargs)
+    return self._modules[name]
+
+
+def var_nest_names(nest):
+  if isinstance(nest, dict):
+    items = ' '.join(f'{k}:{var_nest_names(v)}' for k, v in nest.items())
+    return '{' + items + '}'
+  if isinstance(nest, (list, tuple)):
+    items = ' '.join(var_nest_names(v) for v in nest)
+    return '[' + items + ']'
+  if hasattr(nest, 'name') and hasattr(nest, 'shape'):
+    return nest.name + str(nest.shape).replace(', ', 'x')
+  if hasattr(nest, 'shape'):
+    return str(nest.shape).replace(', ', 'x')
+  return '?'
+
+
+class Logger:
+
+  def __init__(self, logdir, step):
+    self._logdir = logdir
+    self._writer = tf.summary.create_file_writer(str(logdir), max_queue=1000)
+    self._last_step = None
+    self._last_time = None
+    self._scalars = {}
+    self._images = {}
+    self._videos = {}
+    self.step = step
+
+  def scalar(self, name, value):
+    self._scalars[name] = float(value)
+
+  def image(self, name, value):
+    self._images[name] = np.array(value)
+
+  def video(self, name, value):
+    self._videos[name] = np.array(value)
+
+  def write(self, fps=False):
+    scalars = list(self._scalars.items())
+    if fps:
+      scalars.append(('fps', self._compute_fps(self.step)))
+    print(f'[{self.step}]', ' / '.join(f'{k} {v:.1f}' for k, v in scalars))
+    with (self._logdir / 'metrics.jsonl').open('a') as f:
+      f.write(json.dumps({'step': self.step, ** dict(scalars)}) + '\n')
+    with self._writer.as_default():
+      for name, value in scalars:
+        tf.summary.scalar('scalars/' + name, value, self.step)
+      for name, value in self._images.items():
+        tf.summary.image(name, value, self.step)
+      for name, value in self._videos.items():
+        video_summary(name, value, self.step)
+    self._writer.flush()
+    self._scalars = {}
+    self._images = {}
+    self._videos = {}
+
+  def _compute_fps(self, step):
+    if self._last_step is None:
+      self._last_time = time.time()
+      self._last_step = step
+      return 0
+    steps = step - self._last_step
+    duration = time.time() - self._last_time
+    self._last_time += duration
+    self._last_step = step
+    return steps / duration
+
+
+def graph_summary(writer, step, fn, *args):
+  def inner(*args):
+    tf.summary.experimental.set_step(step.numpy().item())
+    with writer.as_default():
+      fn(*args)
+  return tf.numpy_function(inner, args, [])
+
+
+def video_summary(name, video, step=None, fps=20):
+  name = name if isinstance(name, str) else name.decode('utf-8')
+  if np.issubdtype(video.dtype, np.floating):
+    video = np.clip(255 * video, 0, 255).astype(np.uint8)
+  B, T, H, W, C = video.shape
+  try:
+    frames = video.transpose((1, 2, 0, 3, 4)).reshape((T, H, B * W, C))
+    summary = tf1.Summary()
+    image = tf1.Summary.Image(height=B * H, width=T * W, colorspace=C)
+    image.encoded_image_string = encode_gif(frames, fps)
+    summary.value.add(tag=name, image=image)
+    tf.summary.experimental.write_raw_pb(summary.SerializeToString(), step)
+  except (IOError, OSError) as e:
+    print('GIF summaries require ffmpeg in $PATH.', e)
+    frames = video.transpose((0, 2, 1, 3, 4)).reshape((1, B * H, T * W, C))
+    tf.summary.image(name, frames, step)
+
+
+def encode_gif(frames, fps):
+  from subprocess import Popen, PIPE
+  h, w, c = frames[0].shape
+  pxfmt = {1: 'gray', 3: 'rgb24'}[c]
+  cmd = ' '.join([
+      f'ffmpeg -y -f rawvideo -vcodec rawvideo',
+      f'-r {fps:.02f} -s {w}x{h} -pix_fmt {pxfmt} -i - -filter_complex',
+      f'[0:v]split[x][z];[z]palettegen[y];[x]fifo[x];[x][y]paletteuse',
+      f'-r {fps:.02f} -f gif -'])
+  proc = Popen(cmd.split(' '), stdin=PIPE, stdout=PIPE, stderr=PIPE)
+  for image in frames:
+    proc.stdin.write(image.tostring())
+  out, err = proc.communicate()
+  if proc.returncode:
+    raise IOError('\n'.join([' '.join(cmd), err.decode('utf8')]))
+  del proc
+  return out
+
+
+def simulate(agent, envs, steps=0, episodes=0, state=None):
+  # Initialize or unpack simulation state.
+  if state is None:
+    step, episode = 0, 0
+    done = np.ones(len(envs), np.bool)
+    length = np.zeros(len(envs), np.int32)
+    obs = [None] * len(envs)
+    agent_state = None
+  else:
+    step, episode, done, length, obs, agent_state = state
+  while (steps and step < steps) or (episodes and episode < episodes):
+    # Reset envs if necessary.
+    if done.any():
+      indices = [index for index, d in enumerate(done) if d]
+      results = [envs[i].reset() for i in indices]
+      for index, result in zip(indices, results):
+        obs[index] = result
+    # Step agents.
+    obs = {k: np.stack([o[k] for o in obs]) for k in obs[0]}
+    action, agent_state = agent(obs, done, agent_state)
+    if isinstance(action, dict):
+      action = [
+          {k: np.array(action[k][i]) for k in action}
+          for i in range(len(envs))]
+    else:
+      action = np.array(action)
+    assert len(action) == len(envs)
+    # Step envs.
+    results = [e.step(a) for e, a in zip(envs, action)]
+    obs, _, done = zip(*[p[:3] for p in results])
+    obs = list(obs)
+    done = np.stack(done)
+    episode += int(done.sum())
+    length += 1
+    step += (done * length).sum()
+    length *= (1 - done)
+    # import pdb
+    # pdb.set_trace()
+  # Return new state to allow resuming the simulation.
+  return (step - steps, episode - episodes, done, length, obs, agent_state)
+
+
+def save_episodes(directory, episodes):
+  directory = pathlib.Path(directory).expanduser()
+  directory.mkdir(parents=True, exist_ok=True)
+  timestamp = datetime.datetime.now().strftime('%Y%m%dT%H%M%S')
+  filenames = []
+  for episode in episodes:
+    identifier = str(uuid.uuid4().hex)
+    length = len(episode['reward'])
+    filename = directory / f'{timestamp}-{identifier}-{length}.npz'
+    with io.BytesIO() as f1:
+      np.savez_compressed(f1, **episode)
+      f1.seek(0)
+      with filename.open('wb') as f2:
+        f2.write(f1.read())
+    filenames.append(filename)
+  return filenames
+
+
+def sample_episodes(episodes, length=None, balance=False, seed=0):
+  random = np.random.RandomState(seed)
+  while True:
+    episode = random.choice(list(episodes.values()))
+    if length:
+      total = len(next(iter(episode.values())))
+      available = total - length
+      if available < 1:
+        # print(f'Skipped short episode of length {available}.')
+        continue
+      if balance:
+        index = min(random.randint(0, total), available)
+      else:
+        index = int(random.randint(0, available + 1))
+      episode = {k: v[index: index + length] for k, v in episode.items()}
+    yield episode
+
+
+def load_episodes(directory, limit=None):
+  directory = pathlib.Path(directory).expanduser()
+  episodes = {}
+  total = 0
+  for filename in reversed(sorted(directory.glob('*.npz'))):
+    try:
+      with filename.open('rb') as f:
+        episode = np.load(f)
+        episode = {k: episode[k] for k in episode.keys()}
+    except Exception as e:
+      print(f'Could not load episode: {e}')
+      continue
+    episodes[str(filename)] = episode
+    total += len(episode['reward']) - 1
+    if limit and total >= limit:
+      break
+  return episodes
+
+
+class DtypeDist:
+
+  def __init__(self, dist, dtype=None):
+    self._dist = dist
+    self._dtype = dtype or prec.global_policy().compute_dtype
+
+  @property
+  def name(self):
+    return 'DtypeDist'
+
+  def __getattr__(self, name):
+    return getattr(self._dist, name)
+
+  def mean(self):
+    return tf.cast(self._dist.mean(), self._dtype)
+
+  def mode(self):
+    return tf.cast(self._dist.mode(), self._dtype)
+
+  def entropy(self):
+    return tf.cast(self._dist.entropy(), self._dtype)
+
+  def sample(self, *args, **kwargs):
+    return tf.cast(self._dist.sample(*args, **kwargs), self._dtype)
+
+
+class SampleDist:
+
+  def __init__(self, dist, samples=100):
+    self._dist = dist
+    self._samples = samples
+
+  @property
+  def name(self):
+    return 'SampleDist'
+
+  def __getattr__(self, name):
+    return getattr(self._dist, name)
+
+  def mean(self):
+    samples = self._dist.sample(self._samples)
+    return tf.reduce_mean(samples, 0)
+
+  def mode(self):
+    sample = self._dist.sample(self._samples)
+    logprob = self._dist.log_prob(sample)
+    return tf.gather(sample, tf.argmax(logprob))[0]
+
+  def entropy(self):
+    sample = self._dist.sample(self._samples)
+    logprob = self.log_prob(sample)
+    return -tf.reduce_mean(logprob, 0)
+
+
+class OneHotDist(tfd.OneHotCategorical):
+
+  def __init__(self, logits=None, probs=None, dtype=None):
+    self._sample_dtype = dtype or prec.global_policy().compute_dtype
+    super().__init__(logits=logits, probs=probs)
+
+  def mode(self):
+    return tf.cast(super().mode(), self._sample_dtype)
+
+  def sample(self, sample_shape=(), seed=None):
+    # Straight through biased gradient estimator.
+    sample = tf.cast(super().sample(sample_shape, seed), self._sample_dtype)
+    probs = super().probs_parameter()
+    while len(probs.shape) < len(sample.shape):
+      probs = probs[None]
+    sample += tf.cast(probs - tf.stop_gradient(probs), self._sample_dtype)
+    return sample
+
+
+class GumbleDist(tfd.RelaxedOneHotCategorical):
+
+  def __init__(self, temp, logits=None, probs=None, dtype=None):
+    self._sample_dtype = dtype or prec.global_policy().compute_dtype
+    self._exact = tfd.OneHotCategorical(logits=logits, probs=probs)
+    super().__init__(temp, logits=logits, probs=probs)
+
+  def mode(self):
+    return tf.cast(self._exact.mode(), self._sample_dtype)
+
+  def entropy(self):
+    return tf.cast(self._exact.entropy(), self._sample_dtype)
+
+  def sample(self, sample_shape=(), seed=None):
+    return tf.cast(super().sample(sample_shape, seed), self._sample_dtype)
+
+
+class UnnormalizedHuber(tfd.Normal):
+
+  def __init__(self, loc, scale, threshold=1, **kwargs):
+    self._threshold = tf.cast(threshold, loc.dtype)
+    super().__init__(loc, scale, **kwargs)
+
+  def log_prob(self, event):
+    return -(tf.math.sqrt(
+        (event - self.mean()) ** 2 + self._threshold ** 2) - self._threshold)
+
+
+class SafeTruncatedNormal(tfd.TruncatedNormal):
+
+  def __init__(self, loc, scale, low, high, clip=1e-6, mult=1):
+    super().__init__(loc, scale, low, high)
+    self._clip = clip
+    self._mult = mult
+
+  def sample(self, *args, **kwargs):
+    event = super().sample(*args, **kwargs)
+    if self._clip:
+      clipped = tf.clip_by_value(
+          event, self.low + self._clip, self.high - self._clip)
+      event = event - tf.stop_gradient(event) + tf.stop_gradient(clipped)
+    if self._mult:
+      event *= self._mult
+    return event
+
+
+class TanhBijector(tfp.bijectors.Bijector):
+
+  def __init__(self, validate_args=False, name='tanh'):
+    super().__init__(
+        forward_min_event_ndims=0,
+        validate_args=validate_args,
+        name=name)
+
+  def _forward(self, x):
+    return tf.nn.tanh(x)
+
+  def _inverse(self, y):
+    dtype = y.dtype
+    y = tf.cast(y, tf.float32)
+    y = tf.where(
+        tf.less_equal(tf.abs(y), 1.),
+        tf.clip_by_value(y, -0.99999997, 0.99999997), y)
+    y = tf.atanh(y)
+    y = tf.cast(y, dtype)
+    return y
+
+  def _forward_log_det_jacobian(self, x):
+    log2 = tf.math.log(tf.constant(2.0, dtype=x.dtype))
+    return 2.0 * (log2 - x - tf.nn.softplus(-2.0 * x))
+
+
+def lambda_return(
+    reward, value, pcont, bootstrap, lambda_, axis):
+  # Setting lambda=1 gives a discounted Monte Carlo return.
+  # Setting lambda=0 gives a fixed 1-step return.
+  assert reward.shape.ndims == value.shape.ndims, (reward.shape, value.shape)
+  if isinstance(pcont, (int, float)):
+    pcont = pcont * tf.ones_like(reward)
+  dims = list(range(reward.shape.ndims))
+  dims = [axis] + dims[1:axis] + [0] + dims[axis + 1:]
+  if axis != 0:
+    reward = tf.transpose(reward, dims)
+    value = tf.transpose(value, dims)
+    pcont = tf.transpose(pcont, dims)
+  if bootstrap is None:
+    bootstrap = tf.zeros_like(value[-1])
+  next_values = tf.concat([value[1:], bootstrap[None]], 0)
+  inputs = reward + pcont * next_values * (1 - lambda_)
+  returns = static_scan(
+      lambda agg, cur: cur[0] + cur[1] * lambda_ * agg,
+      (inputs, pcont), bootstrap, reverse=True)
+  if axis != 0:
+    returns = tf.transpose(returns, dims)
+  return returns
+
+
+class Optimizer(tf.Module):
+
+  def __init__(
+      self, name, lr, eps=1e-4, clip=None, wd=None, wd_pattern=r'.*',
+      opt='adam'):
+    assert 0 <= wd < 1
+    assert not clip or 1 <= clip
+    self._name = name
+    self._clip = clip
+    self._wd = wd
+    self._wd_pattern = wd_pattern
+    self._opt = {
+        'adam': lambda: tf.optimizers.Adam(lr, epsilon=eps),
+        'nadam': lambda: tf.optimizers.Nadam(lr, epsilon=eps),
+        'adamax': lambda: tf.optimizers.Adamax(lr, epsilon=eps),
+        'sgd': lambda: tf.optimizers.SGD(lr),
+        'momentum': lambda: tf.optimizers.SGD(lr, 0.9),
+    }[opt]()
+    self._mixed = (prec.global_policy().compute_dtype == tf.float16)
+    if self._mixed:
+      self._opt = prec.LossScaleOptimizer(self._opt, 'dynamic')
+
+  @property
+  def variables(self):
+    return self._opt.variables()
+
+  def __call__(self, tape, loss, modules, prefix=None):
+    assert loss.dtype is tf.float32, self._name
+    modules = modules if hasattr(modules, '__len__') else (modules,)
+    varibs = tf.nest.flatten([module.variables for module in modules])
+    count = sum(np.prod(x.shape) for x in varibs)
+    print(f'Found {count} {self._name} parameters.')
+    assert len(loss.shape) == 0, loss.shape
+    tf.debugging.check_numerics(loss, self._name + '_loss')
+    if self._mixed:
+      with tape:
+        loss = self._opt.get_scaled_loss(loss)
+    grads = tape.gradient(loss, varibs)
+    if self._mixed:
+      grads = self._opt.get_unscaled_gradients(grads)
+    norm = tf.linalg.global_norm(grads)
+    if not self._mixed:
+      tf.debugging.check_numerics(norm, self._name + '_norm')
+    if self._clip:
+      grads, _ = tf.clip_by_global_norm(grads, self._clip, norm)
+    if self._wd:
+      self._apply_weight_decay(varibs)
+    self._opt.apply_gradients(zip(grads, varibs))
+    metrics = {}
+    if prefix:
+      metrics[f'{prefix}/{self._name}_loss'] = loss
+      metrics[f'{prefix}/{self._name}_grad_norm'] = norm
+      if self._mixed:
+        metrics[f'{prefix}/{self._name}_loss_scale'] = \
+            self._opt.loss_scale._current_loss_scale
+    else:
+      metrics[f'{self._name}_loss'] = loss
+      metrics[f'{self._name}_grad_norm'] = norm
+      if self._mixed:
+        metrics[f'{self._name}_loss_scale'] = \
+            self._opt.loss_scale._current_loss_scale
+    return metrics
+
+  def _apply_weight_decay(self, varibs):
+    nontrivial = (self._wd_pattern != r'.*')
+    if nontrivial:
+      print('Applied weight decay to variables:')
+    for var in varibs:
+      if re.search(self._wd_pattern, self._name + '/' + var.name):
+        if nontrivial:
+          print('- ' + self._name + '/' + var.name)
+        var.assign((1 - self._wd) * var)
+
+
+def args_type(default):
+  def parse_string(x):
+    if default is None:
+      return x
+    if isinstance(default, bool):
+      return bool(['False', 'True'].index(x))
+    if isinstance(default, int):
+      return float(x) if ('e' in x or '.' in x) else int(x)
+    if isinstance(default, (list, tuple)):
+      return tuple(args_type(default[0])(y) for y in x.split(','))
+    return type(default)(x)
+  def parse_object(x):
+    if isinstance(default, (list, tuple)):
+      return tuple(x)
+    return x
+  return lambda x: parse_string(x) if isinstance(x, str) else parse_object(x)
+
+
+def static_scan(fn, inputs, start, reverse=False):
+  last = start
+  outputs = [[] for _ in tf.nest.flatten(start)]
+  indices = range(len(tf.nest.flatten(inputs)[0]))
+  if reverse:
+    indices = reversed(indices)
+  for index in indices:
+    inp = tf.nest.map_structure(lambda x: x[index], inputs)
+    last = fn(last, inp)
+    [o.append(l) for o, l in zip(outputs, tf.nest.flatten(last))]
+  if reverse:
+    outputs = [list(reversed(x)) for x in outputs]
+  outputs = [tf.stack(x, 0) for x in outputs]
+  return tf.nest.pack_sequence_as(start, outputs)
+
+
+def uniform_mixture(dist, dtype=None):
+  if dist.batch_shape[-1] == 1:
+    return tfd.BatchReshape(dist, dist.batch_shape[:-1])
+  dtype = dtype or prec.global_policy().compute_dtype
+  weights = tfd.Categorical(tf.zeros(dist.batch_shape, dtype))
+  return tfd.MixtureSameFamily(weights, dist)
+
+
+def cat_mixture_entropy(dist):
+  if isinstance(dist, tfd.MixtureSameFamily):
+    probs = dist.components_distribution.probs_parameter()
+  else:
+    probs = dist.probs_parameter()
+  return -tf.reduce_mean(
+      tf.reduce_mean(probs, 2) *
+      tf.math.log(tf.reduce_mean(probs, 2) + 1e-8), -1)
+
+
+@tf.function
+def cem_planner(
+    state, num_actions, horizon, proposals, topk, iterations, imagine,
+    objective):
+  dtype = prec.global_policy().compute_dtype
+  B, P = list(state.values())[0].shape[0], proposals
+  H, A = horizon, num_actions
+  flat_state = {k: tf.repeat(v, P, 0) for k, v in state.items()}
+  mean = tf.zeros((B, H, A), dtype)
+  std = tf.ones((B, H, A), dtype)
+  for _ in range(iterations):
+    proposals = tf.random.normal((B, P, H, A), dtype=dtype)
+    proposals = proposals * std[:, None] + mean[:, None]
+    proposals = tf.clip_by_value(proposals, -1, 1)
+    flat_proposals = tf.reshape(proposals, (B * P, H, A))
+    states = imagine(flat_proposals, flat_state)
+    scores = objective(states)
+    scores = tf.reshape(tf.reduce_sum(scores, -1), (B, P))
+    _, indices = tf.math.top_k(scores, topk, sorted=False)
+    best = tf.gather(proposals, indices, axis=1, batch_dims=1)
+    mean, var = tf.nn.moments(best, 1)
+    std = tf.sqrt(var + 1e-6)
+  return mean[:, 0, :]
+
+
+@tf.function
+def grad_planner(
+    state, num_actions, horizon, proposals, iterations, imagine, objective,
+    kl_scale, step_size):
+  dtype = prec.global_policy().compute_dtype
+  B, P = list(state.values())[0].shape[0], proposals
+  H, A = horizon, num_actions
+  flat_state = {k: tf.repeat(v, P, 0) for k, v in state.items()}
+  mean = tf.zeros((B, H, A), dtype)
+  rawstd = 0.54 * tf.ones((B, H, A), dtype)
+  for _ in range(iterations):
+    proposals = tf.random.normal((B, P, H, A), dtype=dtype)
+    with tf.GradientTape(watch_accessed_variables=False) as tape:
+      tape.watch(mean)
+      tape.watch(rawstd)
+      std = tf.nn.softplus(rawstd)
+      proposals = proposals * std[:, None] + mean[:, None]
+      proposals = (
+          tf.stop_gradient(tf.clip_by_value(proposals, -1, 1)) +
+          proposals - tf.stop_gradient(proposals))
+      flat_proposals = tf.reshape(proposals, (B * P, H, A))
+      states = imagine(flat_proposals, flat_state)
+      scores = objective(states)
+      scores = tf.reshape(tf.reduce_sum(scores, -1), (B, P))
+      div = tfd.kl_divergence(
+          tfd.Normal(mean, std),
+          tfd.Normal(tf.zeros_like(mean), tf.ones_like(std)))
+      elbo = tf.reduce_sum(scores) - kl_scale * div
+      elbo /= tf.cast(tf.reduce_prod(tf.shape(scores)), dtype)
+    grad_mean, grad_rawstd = tape.gradient(elbo, [mean, rawstd])
+    e, v = tf.nn.moments(grad_mean, [1, 2], keepdims=True)
+    grad_mean /= tf.sqrt(e * e + v + 1e-4)
+    e, v = tf.nn.moments(grad_rawstd, [1, 2], keepdims=True)
+    grad_rawstd /= tf.sqrt(e * e + v + 1e-4)
+    mean = tf.clip_by_value(mean + step_size * grad_mean, -1, 1)
+    rawstd = rawstd + step_size * grad_rawstd
+  return mean[:, 0, :]
+
+
+class Every:
+
+  def __init__(self, every):
+    self._every = every
+    self._last = None
+
+  def __call__(self, step):
+    if not self._every:
+      return False
+    if self._last is None:
+      self._last = step
+      return True
+    if step >= self._last + self._every:
+      self._last += self._every
+      return True
+    return False
+
+
+class Once:
+
+  def __init__(self):
+    self._once = True
+
+  def __call__(self):
+    if self._once:
+      self._once = False
+      return True
+    return False
+
+
+class Until:
+
+  def __init__(self, until):
+    self._until = until
+
+  def __call__(self, step):
+    if not self._until:
+      return True
+    return step < self._until
+
+
+def schedule(string, step):
+  try:
+    return float(string)
+  except ValueError:
+    step = tf.cast(step, tf.float32)
+    match = re.match(r'linear\((.+),(.+),(.+)\)', string)
+    if match:
+      initial, final, duration = [float(group) for group in match.groups()]
+      mix = tf.clip_by_value(step / duration, 0, 1)
+      return (1 - mix) * initial + mix * final
+    match = re.match(r'warmup\((.+),(.+)\)', string)
+    if match:
+      warmup, value = [float(group) for group in match.groups()]
+      scale = tf.clip_by_value(step / warmup, 0, 1)
+      return scale * value
+    match = re.match(r'exp\((.+),(.+),(.+)\)', string)
+    if match:
+      initial, final, halflife = [float(group) for group in match.groups()]
+      return (initial - final) * 0.5 ** (step / halflife) + final
+    raise NotImplementedError(string)

DreamerV2/wrappers.py

@@ -0,0 +1,280 @@
+import threading
+
+import gym
+import numpy as np
+
+class DeepMindControl:
+
+  def __init__(self, name, action_repeat=1, size=(64, 64), camera=None):
+    domain, task = name.split('_', 1)
+    if domain == 'cup':  # Only domain with multiple words.
+      domain = 'ball_in_cup'
+    if isinstance(domain, str):
+      from dm_control import suite
+      self._env = suite.load(domain, task)
+    else:
+      assert task is None
+      self._env = domain()
+    self._action_repeat = action_repeat
+    self._size = size
+    if camera is None:
+      camera = dict(quadruped=2).get(domain, 0)
+    self._camera = camera
+
+  @property
+  def observation_space(self):
+    spaces = {}
+    for key, value in self._env.observation_spec().items():
+      spaces[key] = gym.spaces.Box(
+          -np.inf, np.inf, value.shape, dtype=np.float32)
+    spaces['image'] = gym.spaces.Box(
+        0, 255, self._size + (3,), dtype=np.uint8)
+    return gym.spaces.Dict(spaces)
+
+  @property
+  def action_space(self):
+    spec = self._env.action_spec()
+    return gym.spaces.Box(spec.minimum, spec.maximum, dtype=np.float32)
+
+  def step(self, action):
+    assert np.isfinite(action).all(), action
+    reward = 0
+    for _ in range(self._action_repeat):
+      time_step = self._env.step(action)
+      reward += time_step.reward or 0
+      if time_step.last():
+        break
+    obs = dict(time_step.observation)
+    obs['image'] = self.render()
+    done = time_step.last()
+    info = {'discount': np.array(time_step.discount, np.float32)}
+    return obs, reward, done, info
+
+  def reset(self):
+    time_step = self._env.reset()
+    obs = dict(time_step.observation)
+    obs['image'] = self.render()
+    return obs
+
+  def render(self, *args, **kwargs):
+    if kwargs.get('mode', 'rgb_array') != 'rgb_array':
+      raise ValueError("Only render mode 'rgb_array' is supported.")
+    return self._env.physics.render(*self._size, camera_id=self._camera)
+
+
+class Atari:
+
+  LOCK = threading.Lock()
+
+  def __init__(
+      self, name, action_repeat=4, size=(84, 84), grayscale=True, noops=30,
+      life_done=False, sticky_actions=True, all_actions=False):
+    assert size[0] == size[1]
+    import gym.wrappers
+    import gym.envs.atari
+    with self.LOCK:
+      env = gym.envs.atari.AtariEnv(
+          game=name, obs_type='image', frameskip=1,
+          repeat_action_probability=0.25 if sticky_actions else 0.0,
+          full_action_space=all_actions)
+    # Avoid unnecessary rendering in inner env.
+    env._get_obs = lambda: None
+    # Tell wrapper that the inner env has no action repeat.
+    env.spec = gym.envs.registration.EnvSpec('NoFrameskip-v0')
+    env = gym.wrappers.AtariPreprocessing(
+        env, noops, action_repeat, size[0], life_done, grayscale)
+    self._env = env
+    self._grayscale = grayscale
+
+  @property
+  def observation_space(self):
+    return gym.spaces.Dict({
+        'image': self._env.observation_space,
+        'ram': gym.spaces.Box(0, 255, (128,), np.uint8),
+    })
+
+  @property
+  def action_space(self):
+    return self._env.action_space
+
+  def close(self):
+    return self._env.close()
+
+  def reset(self):
+    with self.LOCK:
+      image = self._env.reset()
+    if self._grayscale:
+      image = image[..., None]
+    obs = {'image': image, 'ram': self._env.env._get_ram()}
+    return obs
+
+  def step(self, action):
+    image, reward, done, info = self._env.step(action)
+    if self._grayscale:
+      image = image[..., None]
+    obs = {'image': image, 'ram': self._env.env._get_ram()}
+    return obs, reward, done, info
+
+  def render(self, mode):
+    return self._env.render(mode)
+
+class CollectDataset:
+
+  def __init__(self, env, callbacks=None, precision=32):
+    self._env = env
+    self._callbacks = callbacks or ()
+    self._precision = precision
+    self._episode = None
+
+  def __getattr__(self, name):
+    return getattr(self._env, name)
+
+  def step(self, action):
+    obs, reward, done, info = self._env.step(action)
+    obs = {k: self._convert(v) for k, v in obs.items()}
+    transition = obs.copy()
+    transition['action'] = action
+    transition['reward'] = reward
+    transition['discount'] = info.get('discount', np.array(1 - float(done)))
+    self._episode.append(transition)
+    if done:
+      episode = {k: [t[k] for t in self._episode] for k in self._episode[0]}
+      episode = {k: self._convert(v) for k, v in episode.items()}
+      info['episode'] = episode
+      for callback in self._callbacks:
+        callback(episode)
+    return obs, reward, done, info
+
+  def reset(self):
+    obs = self._env.reset()
+    transition = obs.copy()
+    transition['action'] = np.zeros(self._env.action_space.shape)
+    transition['reward'] = 0.0
+    transition['discount'] = 1.0
+    self._episode = [transition]
+    return obs
+
+  def _convert(self, value):
+    value = np.array(value)
+    if np.issubdtype(value.dtype, np.floating):
+      dtype = {16: np.float16, 32: np.float32, 64: np.float64}[self._precision]
+    elif np.issubdtype(value.dtype, np.signedinteger):
+      dtype = {16: np.int16, 32: np.int32, 64: np.int64}[self._precision]
+    elif np.issubdtype(value.dtype, np.uint8):
+      dtype = np.uint8
+    else:
+      raise NotImplementedError(value.dtype)
+    return value.astype(dtype)
+
+
+class TimeLimit:
+
+  def __init__(self, env, duration):
+    self._env = env
+    self._duration = duration
+    self._step = None
+
+  def __getattr__(self, name):
+    return getattr(self._env, name)
+
+  def step(self, action):
+    assert self._step is not None, 'Must reset environment.'
+    obs, reward, done, info = self._env.step(action)
+    self._step += 1
+    if self._step >= self._duration:
+      done = True
+      if 'discount' not in info:
+        info['discount'] = np.array(1.0).astype(np.float32)
+      self._step = None
+    return obs, reward, done, info
+
+  def reset(self):
+    self._step = 0
+    return self._env.reset()
+
+
+class NormalizeActions:
+
+  def __init__(self, env):
+    self._env = env
+    self._mask = np.logical_and(
+        np.isfinite(env.action_space.low),
+        np.isfinite(env.action_space.high))
+    self._low = np.where(self._mask, env.action_space.low, -1)
+    self._high = np.where(self._mask, env.action_space.high, 1)
+
+  def __getattr__(self, name):
+    return getattr(self._env, name)
+
+  @property
+  def action_space(self):
+    low = np.where(self._mask, -np.ones_like(self._low), self._low)
+    high = np.where(self._mask, np.ones_like(self._low), self._high)
+    return gym.spaces.Box(low, high, dtype=np.float32)
+
+  def step(self, action):
+    original = (action + 1) / 2 * (self._high - self._low) + self._low
+    original = np.where(self._mask, original, action)
+    return self._env.step(original)
+
+
+class OneHotAction:
+
+  def __init__(self, env):
+    assert isinstance(env.action_space, gym.spaces.Discrete)
+    self._env = env
+    self._random = np.random.RandomState()
+
+  def __getattr__(self, name):
+    return getattr(self._env, name)
+
+  @property
+  def action_space(self):
+    shape = (self._env.action_space.n,)
+    space = gym.spaces.Box(low=0, high=1, shape=shape, dtype=np.float32)
+    space.sample = self._sample_action
+    return space
+
+  def step(self, action):
+    index = np.argmax(action).astype(int)
+    reference = np.zeros_like(action)
+    reference[index] = 1
+    if not np.allclose(reference, action):
+      raise ValueError(f'Invalid one-hot action:\n{action}')
+    return self._env.step(index)
+
+  def reset(self):
+    return self._env.reset()
+
+  def _sample_action(self):
+    actions = self._env.action_space.n
+    index = self._random.randint(0, actions)
+    reference = np.zeros(actions, dtype=np.float32)
+    reference[index] = 1.0
+    return reference
+
+
+class RewardObs:
+
+  def __init__(self, env):
+    self._env = env
+
+  def __getattr__(self, name):
+    return getattr(self._env, name)
+
+  @property
+  def observation_space(self):
+    spaces = self._env.observation_space.spaces
+    assert 'reward' not in spaces
+    spaces['reward'] = gym.spaces.Box(-np.inf, np.inf, dtype=np.float32)
+    return gym.spaces.Dict(spaces)
+
+  def step(self, action):
+    obs, reward, done, info = self._env.step(action)
+    obs['reward'] = reward
+    return obs, reward, done, info
+
+  def reset(self):
+    obs = self._env.reset()
+    obs['reward'] = 0.0
+    return obs