//  [python]
def learn(
        self,
        n_rollouts=10000,
        n_epochs=15,
        batch_size=1,
        print_evaluation=False,
        play_data=None,
        return_initial_loss=False,
    ):
        """
        Main training loop
        """
        check_greater_equal("n_epochs", n_epochs, 1)
        check_greater_equal("batch_size", batch_size, 1)
        check_greater_equal("n_rollouts", n_rollouts, batch_size)

        use_play_data = play_data is not None
        if use_play_data:
            data = extract_play_data(play_data, self.timesteps)
        else:
            data = self._collect_data(n_rollouts)
        n_samples = len(data[0])
        assert n_samples == len(data[1])

        check_less_equal("batch_size", batch_size, n_samples)
        n_batches = n_samples // batch_size

        if self.sess is None:
            self.sess = get_tf_session()
            self.sess.run(tf.global_variables_initializer())

        if self.tensorboard_log is not None:
            summaries_op = tf.summary.merge_all()
            summary_writer = tf.summary.FileWriter(
                self.tensorboard_log, self.sess.graph
            )
        else:
            summaries_op = tf.no_op()

        first_epoch_losses = []
        last_epoch_losses = []

        for epoch in range(n_epochs):
            for batch in range(n_batches):
                batch_observations, batch_actions = get_batch(
                    data, batch, batch_size, self.timesteps
                )

                batch_loss, batch_elbo, _, batch_lr, summary, step = self.sess.run(
                    [
                        self.loss,
                        self.elbo,
                        self.optimization_op,
                        self.learning_rate,
                        summaries_op,
                        self.global_step,
                    ],
                    feed_dict={
                        self.in_obs_seq: batch_observations,
                        self.in_actions: batch_actions,
                        self.horizon: self.timesteps + 1,
                    },
                )

                if epoch == 0:
                    first_epoch_losses.append(batch_loss)
                if epoch == n_epochs - 1:
                    last_epoch_losses.append(batch_loss)

                if self.tensorboard_log is not None:
                    if (step - 1) % 110 == 0:
                        if use_play_data:
                            self._create_mujoco_vids(
                                self.env, batch_observations, batch_actions, step
                            )
                    summary_writer.add_summary(summary, step)

                print(
                    "Epoch: {}/{}...".format(epoch + 1, n_epochs),
                    "Batch: {}/{}...".format(batch + 1, n_batches),
                    "Training loss: {:.4f}   (learning_rate = {:.6f})".format(
                        batch_loss, batch_lr
                    ),
                    flush=True,
                )

            if self.checkpoint_folder is not None:
                params = {
                    "observation_dist": self.observation_dist,
                    "pixel_observations": self.pixel_observations,
                    "hidden_layer_size": int(self._hidden_layer_size),
                    "n_hidden_layers": int(self._num_layers),
                    "rnn_state_size": int(self._rnn_state_size),
                    "learning_rate": float(self.learning_rate.eval(session=self.sess)),
                    "obs_stddev": float(self._obs_stddev),
                    "likelihood_scale": float(self.likelihood_scale),
                    "mujoco_video_out_path": self.mujoco_video_out_path,
                    "timesteps_training": int(self.timesteps),
                    "fixed_latent_stddev": float(self._fixed_latent_stddev)
                    if self._fixed_latent_stddev is not None
                    else None,
                }
                with open("_".join([self.checkpoint_folder, "params.json"]), "w") as f:
                    json.dump(params, f)
                self.saver.save(self.sess, self.checkpoint_folder)

        if print_evaluation:
            if use_play_data:
                print("Warning: print_evaluation is only supported for gridworlds.")
            else:
                self._print_evaluation()

        if return_initial_loss:
            return np.mean(first_epoch_losses), np.mean(last_epoch_losses)
        return np.mean(last_epoch_losses)