    def _save(self, checkpoint_dir):
        """Implements the checkpoint logic.

        TODO(hartikainen): This implementation is currently very hacky. Things
        that need to be fixed:
          - Figure out how serialize/save tf.keras.Model subclassing. The
            current implementation just dumps the weights in a pickle, which
            is not optimal.
          - Try to unify all the saving and loading into easily
            extendable/maintainable interfaces. Currently we use
            `tf.train.Checkpoint` and `pickle.dump` in very unorganized way
            which makes things not so usable.
        """
        pickle_path = self._pickle_path(checkpoint_dir)
        with open(pickle_path, 'wb') as f:
            pickle.dump(self.picklables, f)

        if self._variant['run_params'].get('checkpoint_replay_pool', False):
            self._save_replay_pool(checkpoint_dir)

        tf_checkpoint = self._get_tf_checkpoint()

        tf_checkpoint.save(
            file_prefix=self._tf_checkpoint_prefix(checkpoint_dir),
            session=self._session)

        return os.path.join(checkpoint_dir, '')
Implements the checkpoint logic.

TODO(hartikainen): This implementation is currently very hacky. Things
that need to be fixed:
- Figure out how serialize/save tf.keras.Model subclassing. The
current implementation just dumps the weights in a pickle, which
is not optimal.
- Try to unify all the saving and loading into easily
extendable/maintainable interfaces. Currently we use
`tf.train.Checkpoint` and `pickle.dump` in very unorganized way
which makes things not so usable.
examples/development/main.py