Change samplers to use random state to allow consistency in reset par… (#2398)

* Change samplers to use random state to allow consistency in reset parameter draws for a specified seed

Author: sankalp04

ml-agents-envs/mlagents/envs/sampler_class.py

@@ -19,13 +19,27 @@ class UniformSampler(Sampler):
     """
 
     def __init__(
-        self, min_value: Union[int, float], max_value: Union[int, float], **kwargs
+        self,
+        min_value: Union[int, float],
+        max_value: Union[int, float],
+        seed: Optional[int] = None,
+        **kwargs
     ) -> None:
+        """
+        :param min_value: minimum value of the range to be sampled uniformly from
+        :param max_value: maximum value of the range to be sampled uniformly from
+        :param seed: Random seed used for making draws from the uniform sampler
+        """
         self.min_value = min_value
         self.max_value = max_value
+        # Draw from random state to allow for consistent reset parameter draw for a seed
+        self.random_state = np.random.RandomState(seed)
 
     def sample_parameter(self) -> float:
-        return np.random.uniform(self.min_value, self.max_value)
+        """
+        Draws and returns a sample from the specified interval
+        """
+        return self.random_state.uniform(self.min_value, self.max_value)
 
 
 class MultiRangeUniformSampler(Sampler):
@@ -36,19 +50,33 @@ class MultiRangeUniformSampler(Sampler):
     it proceeds to pick a value uniformly in that range.
     """
 
-    def __init__(self, intervals: List[List[Union[int, float]]], **kwargs) -> None:
+    def __init__(
+        self,
+        intervals: List[List[Union[int, float]]],
+        seed: Optional[int] = None,
+        **kwargs
+    ) -> None:
+        """
+        :param intervals: List of intervals to draw uniform samples from
+        :param seed: Random seed used for making uniform draws from the specified intervals
+        """
         self.intervals = intervals
         # Measure the length of the intervals
         interval_lengths = [abs(x[1] - x[0]) for x in self.intervals]
         cum_interval_length = sum(interval_lengths)
         # Assign weights to an interval proportionate to the interval size
         self.interval_weights = [x / cum_interval_length for x in interval_lengths]
+        # Draw from random state to allow for consistent reset parameter draw for a seed
+        self.random_state = np.random.RandomState(seed)
 
     def sample_parameter(self) -> float:
+        """
+        Selects an interval to pick and then draws a uniform sample from the picked interval
+        """
         cur_min, cur_max = self.intervals[
-            np.random.choice(len(self.intervals), p=self.interval_weights)
+            self.random_state.choice(len(self.intervals), p=self.interval_weights)
         ]
-        return np.random.uniform(cur_min, cur_max)
+        return self.random_state.uniform(cur_min, cur_max)
 
 
 class GaussianSampler(Sampler):
@@ -58,13 +86,27 @@ class GaussianSampler(Sampler):
     """
 
     def __init__(
-        self, mean: Union[float, int], st_dev: Union[float, int], **kwargs
+        self,
+        mean: Union[float, int],
+        st_dev: Union[float, int],
+        seed: Optional[int] = None,
+        **kwargs
     ) -> None:
+        """
+        :param mean: Specifies the mean of the gaussian distribution to draw from
+        :param st_dev: Specifies the standard devation of the gaussian distribution to draw from
+        :param seed: Random seed used for making gaussian draws from the sample
+        """
         self.mean = mean
         self.st_dev = st_dev
+        # Draw from random state to allow for consistent reset parameter draw for a seed
+        self.random_state = np.random.RandomState(seed)
 
     def sample_parameter(self) -> float:
-        return np.random.normal(self.mean, self.st_dev)
+        """
+        Returns a draw from the specified Gaussian distribution
+        """
+        return self.random_state.normal(self.mean, self.st_dev)
 
 
 class SamplerFactory:
@@ -81,17 +123,31 @@ class SamplerFactory:
 
     @staticmethod
     def register_sampler(name: str, sampler_cls: Type[Sampler]) -> None:
+        """
+        Registers the sampe in the Sampler Factory to be used later
+        :param name: String name to set as key for the sampler_cls in the factory
+        :param sampler_cls: Sampler object to associate to the name in the factory
+        """
         SamplerFactory.NAME_TO_CLASS[name] = sampler_cls
 
     @staticmethod
-    def init_sampler_class(name: str, params: Dict[str, Any]):
+    def init_sampler_class(
+        name: str, params: Dict[str, Any], seed: Optional[int] = None
+    ) -> Sampler:
+        """
+        Initializes the sampler class associated with the name with the params
+        :param name: Name of the sampler in the factory to initialize
+        :param params: Parameters associated to the sampler attached to the name
+        :param seed: Random seed to be used to set deterministic random draws for the sampler
+        """
         if name not in SamplerFactory.NAME_TO_CLASS:
             raise SamplerException(
                 name + " sampler is not registered in the SamplerFactory."
                 " Use the register_sample method to register the string"
                 " associated to your sampler in the SamplerFactory."
             )
         sampler_cls = SamplerFactory.NAME_TO_CLASS[name]
+        params["seed"] = seed
         try:
             return sampler_cls(**params)
         except TypeError:
@@ -103,7 +159,13 @@ def init_sampler_class(name: str, params: Dict[str, Any]):
 
 
 class SamplerManager:
-    def __init__(self, reset_param_dict: Dict[str, Any]) -> None:
+    def __init__(
+        self, reset_param_dict: Dict[str, Any], seed: Optional[int] = None
+    ) -> None:
+        """
+        :param reset_param_dict: Arguments needed for initializing the samplers
+        :param seed: Random seed to be used for drawing samples from the samplers
+        """
         self.reset_param_dict = reset_param_dict if reset_param_dict else {}
         assert isinstance(self.reset_param_dict, dict)
         self.samplers: Dict[str, Sampler] = {}
@@ -116,7 +178,7 @@ def __init__(self, reset_param_dict: Dict[str, Any]) -> None:
                 )
             sampler_name = cur_param_dict.pop("sampler-type")
             param_sampler = SamplerFactory.init_sampler_class(
-                sampler_name, cur_param_dict
+                sampler_name, cur_param_dict, seed
             )
 
             self.samplers[param_name] = param_sampler
@@ -128,6 +190,10 @@ def is_empty(self) -> bool:
         return not bool(self.samplers)
 
     def sample_all(self) -> Dict[str, float]:
+        """
+        Loop over all samplers and draw a sample from each one for generating
+        next set of reset parameter values.
+        """
         res = {}
         for param_name, param_sampler in list(self.samplers.items()):
             res[param_name] = param_sampler.sample_parameter()

ml-agents/mlagents/trainers/learn.py

@@ -91,7 +91,7 @@ def run_training(
     env = SubprocessEnvManager(env_factory, num_envs)
     maybe_meta_curriculum = try_create_meta_curriculum(curriculum_folder, env)
     sampler_manager, resampling_interval = create_sampler_manager(
-        sampler_file_path, env.reset_parameters
+        sampler_file_path, env.reset_parameters, run_seed
     )
 
     # Create controller and begin training.
@@ -118,7 +118,7 @@ def run_training(
     tc.start_learning(env, trainer_config)
 
 
-def create_sampler_manager(sampler_file_path, env_reset_params):
+def create_sampler_manager(sampler_file_path, env_reset_params, run_seed=None):
     sampler_config = None
     resample_interval = None
     if sampler_file_path is not None:
@@ -136,7 +136,7 @@ def create_sampler_manager(sampler_file_path, env_reset_params):
                 "Resampling interval was not specified in the sampler file."
                 " Please specify it with the 'resampling-interval' key in the sampler config file."
             )
-    sampler_manager = SamplerManager(sampler_config)
+    sampler_manager = SamplerManager(sampler_config, run_seed)
     return sampler_manager, resample_interval