# Buggy code
def restore(self, checkpoint_path):
    """Restores training state from a given model checkpoint.

    These checkpoints are returned from calls to save().

    Subclasses should override ``_restore()`` instead to restore state.
    This method restores additional metadata saved with the checkpoint.
    """
    with open(checkpoint_path + ".tune_metadata", "rb") as f:
        metadata = pickle.load(f)
    self._experiment_id = metadata["experiment_id"]
    self._iteration = metadata["iteration"]
    self._timesteps_total = metadata["timesteps_total"]
    self._time_total = metadata["time_total"]
    self._episodes_total = metadata["episodes_total"]
    saved_as_dict = metadata["saved_as_dict"]
    if saved_as_dict:
        with open(checkpoint_path, "rb") as loaded_state:
            checkpoint_dict = pickle.load(loaded_state)
        checkpoint_dict.update(tune_checkpoint_path=checkpoint_path)
        self._restore(checkpoint_dict)
    else:
        self._restore(checkpoint_path)
    self._time_since_restore = 0.0
    self._timesteps_since_restore = 0
    self._iterations_since_restore = 0
    self._restored = True
    logger.info(
        "Restored on %s from checkpoint: %s", self.current_ip(), checkpoint_path
    )
    state = {
        "_iteration": self._iteration,
        "_timesteps_total": self._timesteps_total,
        "_time_total": self._time_total,
        "_episodes_total": self._episodes_total,
    }
    logger.info("Current state after restoring: %s", state)


# Traceback (NotImplementedError)
$ python ./tests/tune_done_test.py
2020-03-22 12:17:21,455 INFO resource_spec.py:212 -- Starting Ray with 10.3 GiB memory available for workers and up to 5.17 GiB for objects. You can adjust these settings with ray.init(memory=<bytes>, object_store_memory=<bytes>).
2020-03-22 12:17:21,746 INFO services.py:1078 -- View the Ray dashboard at localhost:8265
== Status ==
Memory usage on this node: 17.8/32.0 GiB
Using FIFO scheduling algorithm.
Resources requested: 1/16 CPUs, 0/0 GPUs, 0.0/10.3 GiB heap, 0.0/3.56 GiB objects
Result logdir: /Users/hartikainen/ray_results/done-test
Number of trials: 1 (1 RUNNING)
+---------------------------+----------+-------+-------+
| Trial name                | status   | loc   | a/b   |
|---------------------------+----------+-------+-------|
| MyTrainableClass_14fda478 | RUNNING  |       |       |
+---------------------------+----------+-------+-------+


Result for MyTrainableClass_14fda478:
date: 2020-03-22_12-17-23
done: false
episode_reward_mean: 1
experiment_id: 3549c7334c884d3abc309232da4f6679
experiment_tag: '0_b={''d'': ''4''}'
hostname: catz-48fe.stcatz.ox.ac.uk
iterations_since_restore: 1
node_ip: 129.67.48.254
pid: 73820
time_since_restore: 3.0994415283203125e-06
time_this_iter_s: 3.0994415283203125e-06
time_total_s: 3.0994415283203125e-06
timestamp: 1584879443
timesteps_since_restore: 0
training_iteration: 1
trial_id: 14fda478
what:
'1': '2'
'3': 4
'5':
'6': 4

Result for MyTrainableClass_14fda478:
date: 2020-03-22_12-17-23
done: true
episode_reward_mean: 3
experiment_id: 3549c7334c884d3abc309232da4f6679
experiment_tag: '0_b={''d'': ''4''}'
hostname: catz-48fe.stcatz.ox.ac.uk
iterations_since_restore: 3
node_ip: 129.67.48.254
pid: 73820
time_since_restore: 9.775161743164062e-06
time_this_iter_s: 2.86102294921875e-06
time_total_s: 9.775161743164062e-06
timestamp: 1584879443
timesteps_since_restore: 0
training_iteration: 3
trial_id: 14fda478
what:
'1': '2'
'3': 4
'5':
'6': 4

2020-03-22 12:17:23,204 ERROR trial_runner.py:513 -- Trial MyTrainableClass_14fda478: Error processing event.
Traceback (most recent call last):
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/trial_runner.py", line 511, in _process_trial
self._execute_action(trial, decision)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/trial_runner.py", line 595, in _execute_action
self.trial_executor.stop_trial(trial)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/ray_trial_executor.py", line 263, in stop_trial
trial, error=error, error_msg=error_msg, stop_logger=stop_logger)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/ray_trial_executor.py", line 204, in _stop_trial
trial.close_logger()
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/trial.py", line 315, in close_logger
self.result_logger.close()
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/logger.py", line 305, in close
_logger.close()
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/logger.py", line 233, in close
self._try_log_hparams(self.last_result)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/logger.py", line 244, in _try_log_hparams
hparam_dict=scrubbed_params, metric_dict=result)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/tensorboardX/summary.py", line 102, in hparams
v = make_np(v)[0]
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/tensorboardX/x2num.py", line 34, in make_np
'Got {}, but expected numpy array or torch tensor.'.format(type(x)))
NotImplementedError: Got <class 'dict'>, but expected numpy array or torch tensor.
Traceback (most recent call last):
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/trial_runner.py", line 511, in _process_trial
self._execute_action(trial, decision)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/trial_runner.py", line 595, in _execute_action
self.trial_executor.stop_trial(trial)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/ray_trial_executor.py", line 263, in stop_trial
trial, error=error, error_msg=error_msg, stop_logger=stop_logger)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/ray_trial_executor.py", line 204, in _stop_trial
trial.close_logger()
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/trial.py", line 315, in close_logger
self.result_logger.close()
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/logger.py", line 305, in close
_logger.close()
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/logger.py", line 233, in close
self._try_log_hparams(self.last_result)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/logger.py", line 244, in _try_log_hparams
hparam_dict=scrubbed_params, metric_dict=result)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/tensorboardX/summary.py", line 102, in hparams
v = make_np(v)[0]
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/tensorboardX/x2num.py", line 34, in make_np
'Got {}, but expected numpy array or torch tensor.'.format(type(x)))
NotImplementedError: Got <class 'dict'>, but expected numpy array or torch tensor.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "./tests/tune_done_test.py", line 24, in <module>
'b': tune.sample_from(lambda spec: spec['config']['b']['c']),
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/tune.py", line 324, in run
runner.step()
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/trial_runner.py", line 335, in step
self._process_events()  # blocking
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/trial_runner.py", line 444, in _process_events
self._process_trial(trial)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/trial_runner.py", line 514, in _process_trial
self._process_trial_failure(trial, traceback.format_exc())
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/trial_runner.py", line 580, in _process_trial_failure
trial, error=True, error_msg=error_msg)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/ray_trial_executor.py", line 263, in stop_trial
trial, error=error, error_msg=error_msg, stop_logger=stop_logger)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/ray_trial_executor.py", line 204, in _stop_trial
trial.close_logger()
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/trial.py", line 315, in close_logger
self.result_logger.close()
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/logger.py", line 305, in close
_logger.close()
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/logger.py", line 233, in close
self._try_log_hparams(self.last_result)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/ray/tune/logger.py", line 244, in _try_log_hparams
hparam_dict=scrubbed_params, metric_dict=result)
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/tensorboardX/summary.py", line 102, in hparams
v = make_np(v)[0]
File "/Users/hartikainen/conda/envs/softlearning-tf2/lib/python3.7/site-packages/tensorboardX/x2num.py", line 34, in make_np
'Got {}, but expected numpy array or torch tensor.'.format(type(x)))
NotImplementedError: Got <class 'dict'>, but expected numpy array or torch tensor.

# Fixed code
def restore(self, checkpoint_path):
    """Restores training state from a given model checkpoint.

    These checkpoints are returned from calls to save().

    Subclasses should override ``_restore()`` instead to restore state.
    This method restores additional metadata saved with the checkpoint.
    """
    with open(checkpoint_path + ".tune_metadata", "rb") as f:
        metadata = pickle.load(f)
    self._experiment_id = metadata["experiment_id"]
    self._iteration = metadata["iteration"]
    self._timesteps_total = metadata["timesteps_total"]
    self._time_total = metadata["time_total"]
    self._episodes_total = metadata["episodes_total"]
    saved_as_dict = metadata["saved_as_dict"]
    if saved_as_dict:
        with open(checkpoint_path, "rb") as loaded_state:
            checkpoint_dict = pickle.load(loaded_state)
        checkpoint_dict.update(tune_checkpoint_path=checkpoint_path)
        self._restore(checkpoint_dict)
    else:
        self._restore(checkpoint_path)
    self._time_since_restore = 0.0
    self._timesteps_since_restore = 0
    self._iterations_since_restore = 0
    self._restored = True
    logger.info(
        "Restored on %s from checkpoint: %s", self.get_current_ip(), checkpoint_path
    )
    state = {
        "_iteration": self._iteration,
        "_timesteps_total": self._timesteps_total,
        "_time_total": self._time_total,
        "_episodes_total": self._episodes_total,
    }
    logger.info("Current state after restoring: %s", state)