//  [python]
def create_planet_train_fn(model: RecurrentStateSpaceModel = gin.REQUIRED,
    train_steps: int = gin.REQUIRED,
    batch: int = gin.REQUIRED,
    duration: int = gin.REQUIRED,
    learning_rate: float = gin.REQUIRED,
    model_dir=gin.REQUIRED,
    strategy: tf.distribute.Strategy = gin.REQUIRED,
    save_rewards: bool = True):
  """creates a train_fn to train the `tf.Estimator` referenced in state.

  Args:
    model: a reference to the model.
    train_steps: number of training steps.
    batch: the batch size.
    duration: how many timesteps to include in a single video sequence.
    learning_rate: learning rate.
    model_dir: the path to model directory.
    strategy: a tf.distribute.Strategy object.
    save_rewards: whether or not to save the predicted rewards.

  Returns:
    A train_fn with the following positional arguments:
        * data_path: the path to all episodes.
      This function returns nothing.
  """
  iterator = None
  optimizer = tf.keras.optimizers.Adam(learning_rate, epsilon=1e-3)

  @tf.function
  def train_step(obs):

    def train_iter(obs):
      obs = obs.copy()
      obs['image'] = tf.to_float(obs['image']) / 255.0 - 0.5
      with tf.GradientTape() as tape:
        output = model.compute_losses(obs)
        loss, reward_loss, divergence, frames = output[:4]
        pred_rewards, true_rewards, frame_loss = output[4:]
      variables = tape.watched_variables()
      grads = tape.gradient(loss, variables)
      grads, _ = tf.clip_by_global_norm(grads, 1000)
      optimizer.apply_gradients(zip(grads, variables))
      return loss, reward_loss, divergence, frames, pred_rewards, true_rewards, frame_loss

    return strategy.experimental_run_v2(train_iter, args=(obs,))

  def train_fn(data_path):
    """A train_fn to train the planet model."""
    nonlocal iterator
    nonlocal optimizer

    with strategy.scope():
      global_step = tf.Variable(0, dtype=tf.int64, trainable=False)
      checkpoint = tf.train.Checkpoint(
          global_step=global_step,
          optimizer=optimizer,
          **model.get_trackables())
      manager = tf.train.CheckpointManager(checkpoint, model_dir, max_to_keep=1)
      checkpoint.restore(manager.latest_checkpoint)
    if iterator is None:
      dataset = npz.load_dataset_from_directory(data_path, duration, batch)
      dataset = strategy.experimental_distribute_dataset(dataset)
      iterator = dataset

    writer = tfs.create_file_writer(model_dir)
    tfs.experimental.set_step(global_step)
    true_rewards, pred_rewards = None, None
    with writer.as_default():
      for step, obs in enumerate(iterator):
        if step > train_steps:
          if save_rewards:
            # We are only saving the last training batch.
            reward_dir = os.path.join(model_dir, 'train_rewards')
            true_rewards = strategy.experimental_local_results(true_rewards)
            pred_reward = strategy.experimental_local_results(pred_rewards)
            true_rewards = np.concatenate([x.numpy() for x in true_rewards])
            pred_reward = np.concatenate([x.numpy() for x in pred_reward])
            rewards_to_save = {'true': true_rewards, 'pred': pred_reward}
            npz.save_dictionary(rewards_to_save, reward_dir)
          break
        (loss, reward_loss, divergence, frames, pred_rewards, true_rewards,
         frame_loss) = train_step(obs)
        if step % 100 == 0:
          loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, loss)
          reward_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN,
                                        reward_loss)
          divergence = strategy.reduce(tf.distribute.ReduceOp.MEAN, divergence)
          frame_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, frame_loss)
          frames = strategy.experimental_local_results(frames)
          frames = tf.concat(frames, axis=0)
          pred_reward = strategy.experimental_local_results(pred_rewards)
          pred_reward = tf.concat(pred_reward, axis=0)
          tf.logging.info('loss at step %d: %f', step, loss)
          tfs.scalar('loss/total', loss)
          tfs.scalar('loss/reward', reward_loss)
          tfs.scalar('loss/divergence', divergence)
          tfs.scalar('loss/frames', frame_loss)
          tfs.experimental.write_raw_pb(
              visualization.py_gif_summary(
                  tag='predictions/frames',
                  images=frames.numpy(),
                  max_outputs=6,
                  fps=20))
          ground_truth_rewards = (
              tf.concat(
                  strategy.experimental_local_results(obs['reward']),
                  axis=0)[:, :, 0])
          rewards = pred_reward[:, :, 0]
          signals = tf.stack([ground_truth_rewards, rewards], axis=1)
          visualization.py_plot_1d_signal(
              name='predictions/reward',
              signals=signals.numpy(),
              labels=['ground_truth', 'prediction'],
              max_outputs=6)
        global_step.assign_add(1)

      manager.save(global_step)

  return train_fn