Repository: ray-project/ray

Issue:
[tune] PBT should support nested mutations
### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux Ubuntu 16.04 and 18.04
- **Ray installed from (source or binary)**: source@f372f48bf3b51bc4e6b51ad9691f71b4b9004462
- **Ray version**: 0.5.3
- **Python version**: 3.6


I was trying to run PBT scheduler with nested config and mutation. This fails in the [PBTTrialState.explore](https://github.com/ray-project/ray/blob/f372f48bf3b51bc4e6b51ad9691f71b4b9004462/python/ray/tune/schedulers/pbt.py#L49) because it doesn't support nesting. It seems like this should be pretty easily fixable, assuming that there are no further issues with the nesting.

Example config:
```
{'Q_params': {'kwargs': {'hidden_layer_sizes': (256, 256)},
              'type': 'double_feedforward_Q_function'},
 'algorithm_params': {'kwargs': {'action_prior': 'uniform',
                                 'discount': 0.99,
                                 'epoch_length': 1000,
                                 'eval_deterministic': True,
                                 'eval_n_episodes': 1,
                                 'eval_render': False,
                                 'lr': 0.0003,
                                 'n_epochs': 301,
                                 'n_initial_exploration_steps': 10000,
                                 'n_train_repeat': 1,
                                 'reparameterize': True,
                                 'reward_scale': 1.0,
                                 'save_full_state': False,
                                 'store_extra_policy_info': False,
                                 'target_entropy': 'auto',
                                 'target_update_interval': 1,
                                 'tau': 0.005,
                                 'train_every_n_steps': 1},
                      'type': 'SAC'},
 'domain': 'swimmer',
 'env_params': {},
 'git_sha': '8b4f4ab12499b2f10b6cd24fc0112b35cfea9ffd experiment/pbt-test',
 'mode': 'local',
 'policy_params': {'kwargs': {'hidden_layer_sizes': (256, 256),
                              'regularization_coeff': 0.001,
                              'squash': True},
                   'type': 'GaussianPolicy'},
 'preprocessor_params': {'type': None},
 'replay_pool_params': {'kwargs': {'max_size': 1000000.0},
                        'type': 'SimpleReplayPool'},
 'run_params': {'checkpoint_at_end': True,
                'checkpoint_frequency': 60,
                'seed': 2},
 'sampler_params': {'kwargs': {'batch_size': 256,
                               'max_path_length': 1000,
                               'min_pool_size': 1000},
                    'type': 'SimpleSampler'},
 'task': 'default',
 'universe': 'gym'}
```
Example mutation:
```
{'algorithm_params': {'kwargs': {'discount': <function launch_experiments_ray.<locals>.<lambda> at 0x7f21de052950>,
                                 'lr': <function launch_experiments_ray.<locals>.<lambda> at 0x7f21de0529d8>,
                                 'n_train_repeat': [1, 2, 4, 8]}}}
```


Discussion:

Resolving patch:
diff --git a/python/ray/tune/schedulers/pbt.py b/python/ray/tune/schedulers/pbt.py
--- a/python/ray/tune/schedulers/pbt.py
+++ b/python/ray/tune/schedulers/pbt.py
@@ -47,7 +47,12 @@ def explore(config, mutations, resample_probability, custom_explore_fn):
     """
     new_config = copy.deepcopy(config)
     for key, distribution in mutations.items():
-        if isinstance(distribution, list):
+        if isinstance(distribution, dict):
+            new_config.update({
+                key: explore(config[key], mutations[key], resample_probability,
+                             None)
+            })
+        elif isinstance(distribution, list):
             if random.random() < resample_probability or \
                     config[key] not in distribution:
                 new_config[key] = random.choice(distribution)
@@ -213,8 +218,8 @@ def _exploit(self, trial_executor, trial, trial_to_clone):
         trial_state = self._trial_state[trial]
         new_state = self._trial_state[trial_to_clone]
         if not new_state.last_checkpoint:
-            logger.warning("[pbt]: no checkpoint for trial"
-                           "skip exploit for Trial {}".format(trial))
+            logger.warning("[pbt]: no checkpoint for trial."
+                           " Skip exploit for Trial {}".format(trial))
             return
         new_config = explore(trial_to_clone.config, self._hyperparam_mutations,
                              self._resample_probability,


Test patch:
diff --git a/python/ray/tune/test/trial_scheduler_test.py b/python/ray/tune/test/trial_scheduler_test.py
--- a/python/ray/tune/test/trial_scheduler_test.py
+++ b/python/ray/tune/test/trial_scheduler_test.py
@@ -5,7 +5,7 @@
 import random
 import unittest
 import numpy as np
-
+import sys
 import ray
 from ray.tune.schedulers import (HyperBandScheduler, AsyncHyperBandScheduler,
                                  PopulationBasedTraining, MedianStoppingRule,
@@ -17,6 +17,11 @@
 from ray.rllib import _register_all
 _register_all()
 
+if sys.version_info >= (3, 3):
+    from unittest.mock import MagicMock
+else:
+    from mock import MagicMock
+
 
 def result(t, rew):
     return dict(
@@ -747,6 +752,104 @@ def assertProduces(fn, values):
                 lambda x: x),
             {10.0, 100.0})
 
+        def deep_add(seen, new_values):
+            for k, new_value in new_values.items():
+                if isinstance(new_value, dict):
+                    if k not in seen:
+                        seen[k] = {}
+                    seen[k].update(deep_add(seen[k], new_value))
+                else:
+                    if k not in seen:
+                        seen[k] = set()
+                    seen[k].add(new_value)
+
+            return seen
+
+        def assertNestedProduces(fn, values):
+            random.seed(0)
+            seen = {}
+            for _ in range(100):
+                new_config = fn()
+                seen = deep_add(seen, new_config)
+            self.assertEqual(seen, values)
+
+        # Nested mutation and spec
+        assertNestedProduces(
+            lambda: explore(
+                {
+                    "a": {
+                        "b": 4
+                    },
+                    "1": {
+                        "2": {
+                            "3": 100
+                        }
+                    },
+                },
+                {
+                    "a": {
+                        "b": [3, 4, 8, 10]
+                    },
+                    "1": {
+                        "2": {
+                            "3": lambda: random.choice([10, 100])
+                        }
+                    },
+                },
+                0.0,
+                lambda x: x),
+            {
+                "a": {
+                    "b": {3, 8}
+                },
+                "1": {
+                    "2": {
+                        "3": {80, 120}
+                    }
+                },
+            })
+
+        custom_explore_fn = MagicMock(side_effect=lambda x: x)
+
+        # Nested mutation and spec
+        assertNestedProduces(
+            lambda: explore(
+                {
+                    "a": {
+                        "b": 4
+                    },
+                    "1": {
+                        "2": {
+                            "3": 100
+                        }
+                    },
+                },
+                {
+                    "a": {
+                        "b": [3, 4, 8, 10]
+                    },
+                    "1": {
+                        "2": {
+                            "3": lambda: random.choice([10, 100])
+                        }
+                    },
+                },
+                0.0,
+                custom_explore_fn),
+            {
+                "a": {
+                    "b": {3, 8}
+                },
+                "1": {
+                    "2": {
+                        "3": {80, 120}
+                    }
+                },
+            })
+
+        # Expect call count to be 100 because we call explore 100 times
+        self.assertEqual(custom_explore_fn.call_count, 100)
+
     def testYieldsTimeToOtherTrials(self):
         pbt, runner = self.basicSetup()
         trials = runner.get_trials()