diff --git a/rllib/agents/ars/ars.py b/rllib/agents/ars/ars.py index e4b2cc99c..ef74f1a3e 100644 --- a/rllib/agents/ars/ars.py +++ b/rllib/agents/ars/ars.py @@ -9,7 +9,6 @@ import time import ray from ray.rllib.agents import Trainer, with_common_config - from ray.rllib.agents.ars.ars_tf_policy import ARSTFPolicy from ray.rllib.agents.es import optimizers, utils from ray.rllib.agents.es.es import validate_config @@ -40,6 +39,14 @@ DEFAULT_CONFIG = with_common_config({ "eval_prob": 0.03, # probability of evaluating the parameter rewards "report_length": 10, # how many of the last rewards we average over "offset": 0, + # ARS will use Trainer's evaluation WorkerSet (if evaluation_interval > 0). + # Therefore, we must be careful not to use more than 1 env per eval worker + # (would break ARSPolicy's compute_action method) and to not do obs- + # filtering. + "evaluation_config": { + "num_envs_per_worker": 1, + "observation_filter": "NoFilter" + }, }) # __sphinx_doc_end__ # yapf: enable @@ -183,9 +190,9 @@ class ARSTrainer(Trainer): env_context = EnvContext(config["env_config"] or {}, worker_index=0) env = env_creator(env_context) - policy_cls = get_policy_class(config) - self.policy = policy_cls(env.observation_space, env.action_space, - config) + self._policy_class = get_policy_class(config) + self.policy = self._policy_class(env.observation_space, + env.action_space, config) self.optimizer = optimizers.SGD(self.policy, config["sgd_stepsize"]) self.rollouts_used = config["rollouts_used"] @@ -320,10 +327,19 @@ class ARSTrainer(Trainer): @override(Trainer) def compute_action(self, observation, *args, **kwargs): - action = self.policy.compute_actions(observation, update=True)[0] + action, _, _ = self.policy.compute_actions(observation, update=True) if kwargs.get("full_fetch"): - return action, [], {} - return action + return action[0], [], {} + return action[0] + + @override(Trainer) + def _sync_weights_to_workers(self, *, worker_set=None, workers=None): + # Broadcast the new policy weights to all evaluation workers. + assert worker_set is not None + logger.info("Synchronizing weights to evaluation workers.") + weights = ray.put(self.policy.get_flat_weights()) + worker_set.foreach_policy( + lambda p, pid: p.set_flat_weights(ray.get(weights))) def _collect_results(self, theta_id, min_episodes): num_episodes, num_timesteps = 0, 0 diff --git a/rllib/agents/ars/ars_tf_policy.py b/rllib/agents/ars/ars_tf_policy.py index 53760727e..b482449d8 100644 --- a/rllib/agents/ars/ars_tf_policy.py +++ b/rllib/agents/ars/ars_tf_policy.py @@ -9,6 +9,7 @@ import ray import ray.experimental.tf_utils from ray.rllib.agents.es.es_tf_policy import make_session from ray.rllib.models import ModelCatalog +from ray.rllib.policy.policy import Policy from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.utils.filter import get_filter from ray.rllib.utils.framework import try_import_tf @@ -17,18 +18,17 @@ from ray.rllib.utils.spaces.space_utils import unbatch tf1, tf, tfv = try_import_tf() -class ARSTFPolicy: +class ARSTFPolicy(Policy): def __init__(self, obs_space, action_space, config): - self.observation_space = obs_space - self.action_space = action_space - self.action_noise_std = config["action_noise_std"] + super().__init__(obs_space, action_space, config) + self.action_noise_std = self.config["action_noise_std"] self.preprocessor = ModelCatalog.get_preprocessor_for_space( self.observation_space) - self.observation_filter = get_filter(config["observation_filter"], + self.observation_filter = get_filter(self.config["observation_filter"], self.preprocessor.shape) - self.single_threaded = config.get("single_threaded", False) - if config["framework"] == "tf": + self.single_threaded = self.config.get("single_threaded", False) + if self.config["framework"] == "tf": self.sess = make_session(single_threaded=self.single_threaded) self.inputs = tf1.placeholder( tf.float32, [None] + list(self.preprocessor.shape)) @@ -39,13 +39,13 @@ class ARSTFPolicy: # Policy network. self.dist_class, dist_dim = ModelCatalog.get_action_dist( - self.action_space, config["model"], dist_type="deterministic") + self.action_space, self.config["model"], dist_type="deterministic") self.model = ModelCatalog.get_model_v2( obs_space=self.preprocessor.observation_space, action_space=self.action_space, num_outputs=dist_dim, - model_config=config["model"]) + model_config=self.config["model"]) self.sampler = None if self.sess: @@ -89,16 +89,16 @@ class ARSTFPolicy: actions = unbatch(actions) if add_noise and isinstance(self.action_space, gym.spaces.Box): actions += np.random.randn(*actions.shape) * self.action_noise_std - return actions + return actions, [], {} def compute_single_action(self, observation, add_noise=False, update=True, **kwargs): - action = self.compute_actions( + action, state_outs, extra_fetches = self.compute_actions( [observation], add_noise=add_noise, update=update, **kwargs) - return action[0], [], {} + return action[0], state_outs, extra_fetches def get_state(self): return {"state": self.get_flat_weights()} diff --git a/rllib/agents/ars/tests/test_ars.py b/rllib/agents/ars/tests/test_ars.py index 2bf0b5470..b6bb3c8df 100644 --- a/rllib/agents/ars/tests/test_ars.py +++ b/rllib/agents/ars/tests/test_ars.py @@ -9,11 +9,15 @@ from ray.rllib.utils.test_utils import framework_iterator, \ class TestARS(unittest.TestCase): def test_ars_compilation(self): """Test whether an ARSTrainer can be built on all frameworks.""" - ray.init(num_cpus=2, local_mode=True) + ray.init(num_cpus=3) config = ars.DEFAULT_CONFIG.copy() # Keep it simple. config["model"]["fcnet_hiddens"] = [10] config["model"]["fcnet_activation"] = None + config["noise_size"] = 2500000 + # Test eval workers ("normal" Trainer eval WorkerSet, unusual for ARS). + config["evaluation_interval"] = 1 + config["evaluation_num_workers"] = 1 num_iterations = 2 diff --git a/rllib/agents/es/es.py b/rllib/agents/es/es.py index f94fc469a..d669bd5d0 100644 --- a/rllib/agents/es/es.py +++ b/rllib/agents/es/es.py @@ -37,6 +37,14 @@ DEFAULT_CONFIG = with_common_config({ "observation_filter": "MeanStdFilter", "noise_size": 250000000, "report_length": 10, + # ARS will use Trainer's evaluation WorkerSet (if evaluation_interval > 0). + # Therefore, we must be careful not to use more than 1 env per eval worker + # (would break ESPolicy's compute_action method) and to not do obs- + # filtering. + "evaluation_config": { + "num_envs_per_worker": 1, + "observation_filter": "NoFilter" + }, }) # __sphinx_doc_end__ # yapf: enable @@ -83,9 +91,9 @@ class Worker: self.preprocessor = models.ModelCatalog.get_preprocessor( self.env, config["model"]) - policy_cls = get_policy_class(config) - self.policy = policy_cls(self.env.observation_space, - self.env.action_space, config) + _policy_class = get_policy_class(config) + self.policy = _policy_class(self.env.observation_space, + self.env.action_space, config) @property def filters(self): @@ -172,6 +180,15 @@ def get_policy_class(config): def validate_config(config): if config["num_workers"] <= 0: raise ValueError("`num_workers` must be > 0 for ES!") + if config["evaluation_config"]["num_envs_per_worker"] != 1: + raise ValueError( + "`evaluation_config.num_envs_per_worker` must always be 1 for " + "ES/ARS! To parallelize evaluation, increase " + "`evaluation_num_workers` to > 1.") + if config["evaluation_config"]["observation_filter"] != "NoFilter": + raise ValueError( + "`evaluation_config.observation_filter` must always be `NoFilter` " + "for ES/ARS!") class ESTrainer(Trainer): @@ -185,8 +202,8 @@ class ESTrainer(Trainer): validate_config(config) env_context = EnvContext(config["env_config"] or {}, worker_index=0) env = env_creator(env_context) - policy_cls = get_policy_class(config) - self.policy = policy_cls( + self._policy_class = get_policy_class(config) + self.policy = self._policy_class( obs_space=env.observation_space, action_space=env.action_space, config=config) @@ -307,10 +324,19 @@ class ESTrainer(Trainer): @override(Trainer) def compute_action(self, observation, *args, **kwargs): - action = self.policy.compute_actions(observation, update=False)[0] + action, _, _ = self.policy.compute_actions(observation, update=False) if kwargs.get("full_fetch"): - return action, [], {} - return action + return action[0], [], {} + return action[0] + + @override(Trainer) + def _sync_weights_to_workers(self, *, worker_set=None, workers=None): + # Broadcast the new policy weights to all evaluation workers. + assert worker_set is not None + logger.info("Synchronizing weights to evaluation workers.") + weights = ray.put(self.policy.get_flat_weights()) + worker_set.foreach_policy( + lambda p, pid: p.set_flat_weights(ray.get(weights))) @override(Trainer) def cleanup(self): diff --git a/rllib/agents/es/es_tf_policy.py b/rllib/agents/es/es_tf_policy.py index dc6833f17..ad19b7ba5 100644 --- a/rllib/agents/es/es_tf_policy.py +++ b/rllib/agents/es/es_tf_policy.py @@ -8,7 +8,9 @@ import tree import ray import ray.experimental.tf_utils from ray.rllib.models import ModelCatalog +from ray.rllib.policy.policy import Policy from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import override from ray.rllib.utils.filter import get_filter from ray.rllib.utils.framework import try_import_tf from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space, \ @@ -44,8 +46,9 @@ def rollout(policy, env, timestep_limit=None, add_noise=False, offset=0.0): t = 0 observation = env.reset() for _ in range(timestep_limit or max_timestep_limit): - ac = policy.compute_actions( - observation, add_noise=add_noise, update=True)[0] + ac, _, _ = policy.compute_actions( + observation, add_noise=add_noise, update=True) + ac = ac[0] observation, r, done, _ = env.step(ac) if offset != 0.0: r -= np.abs(offset) @@ -65,17 +68,16 @@ def make_session(single_threaded): inter_op_parallelism_threads=1, intra_op_parallelism_threads=1)) -class ESTFPolicy: +class ESTFPolicy(Policy): def __init__(self, obs_space, action_space, config): - self.observation_space = obs_space - self.action_space = action_space + super().__init__(obs_space, action_space, config) self.action_space_struct = get_base_struct_from_space(action_space) - self.action_noise_std = config["action_noise_std"] + self.action_noise_std = self.config["action_noise_std"] self.preprocessor = ModelCatalog.get_preprocessor_for_space(obs_space) - self.observation_filter = get_filter(config["observation_filter"], + self.observation_filter = get_filter(self.config["observation_filter"], self.preprocessor.shape) - self.single_threaded = config.get("single_threaded", False) - if config["framework"] == "tf": + self.single_threaded = self.config.get("single_threaded", False) + if self.config["framework"] == "tf": self.sess = make_session(single_threaded=self.single_threaded) self.inputs = tf1.placeholder( tf.float32, [None] + list(self.preprocessor.shape)) @@ -86,13 +88,13 @@ class ESTFPolicy: # Policy network. self.dist_class, dist_dim = ModelCatalog.get_action_dist( - self.action_space, config["model"], dist_type="deterministic") + self.action_space, self.config["model"], dist_type="deterministic") self.model = ModelCatalog.get_model_v2( obs_space=self.preprocessor.observation_space, action_space=action_space, num_outputs=dist_dim, - model_config=config["model"]) + model_config=self.config["model"]) self.sampler = None if self.sess: @@ -110,6 +112,7 @@ class ESTFPolicy: np.prod(variable.shape.as_list()) for _, variable in self.variables.variables.items()) + @override(Policy) def compute_actions(self, observation, add_noise=False, @@ -138,16 +141,16 @@ class ESTFPolicy: # Convert `flat_actions` to a list of lists of action components # (list of single actions). actions = unbatch(actions) - return actions + return actions, [], {} def compute_single_action(self, observation, add_noise=False, update=True, **kwargs): - action = self.compute_actions( + action, state_outs, extra_fetches = self.compute_actions( [observation], add_noise=add_noise, update=update, **kwargs) - return action[0], [], {} + return action[0], state_outs, extra_fetches def _add_noise(self, single_action, single_action_space): if isinstance(single_action_space, gym.spaces.Box): diff --git a/rllib/agents/es/es_torch_policy.py b/rllib/agents/es/es_torch_policy.py index 6812ebed3..6f7e374c9 100644 --- a/rllib/agents/es/es_torch_policy.py +++ b/rllib/agents/es/es_torch_policy.py @@ -48,7 +48,7 @@ def before_init(policy, observation_space, action_space, config): for k in sorted(theta_dict.keys()): theta_list.append(torch.reshape(theta_dict[k], (-1, ))) cat = torch.cat(theta_list, dim=0) - return cat.numpy() + return cat.cpu().numpy() type(policy).set_flat_weights = _set_flat_weights type(policy).get_flat_weights = _get_flat_weights @@ -73,7 +73,7 @@ def before_init(policy, observation_space, action_space, config): action = dist.sample() def _add_noise(single_action, single_action_space): - single_action = single_action.detach().numpy() + single_action = single_action.detach().cpu().numpy() if add_noise and isinstance(single_action_space, gym.spaces.Box): single_action += np.random.randn(*single_action.shape) * \ policy.action_noise_std @@ -82,9 +82,19 @@ def before_init(policy, observation_space, action_space, config): action = tree.map_structure(_add_noise, action, policy.action_space_struct) action = unbatch(action) - return action + return action, [], {} + + def _compute_single_action(policy, + observation, + add_noise=False, + update=True, + **kwargs): + action, state_outs, extra_fetches = policy.compute_actions( + [observation], add_noise=add_noise, update=update, **kwargs) + return action[0], state_outs, extra_fetches type(policy).compute_actions = _compute_actions + type(policy).compute_single_action = _compute_single_action def after_init(policy, observation_space, action_space, config): diff --git a/rllib/agents/es/tests/test_es.py b/rllib/agents/es/tests/test_es.py index 57b2dc615..c29761193 100644 --- a/rllib/agents/es/tests/test_es.py +++ b/rllib/agents/es/tests/test_es.py @@ -9,7 +9,7 @@ from ray.rllib.utils.test_utils import check_compute_single_action, \ class TestES(unittest.TestCase): def test_es_compilation(self): """Test whether an ESTrainer can be built on all frameworks.""" - ray.init(num_cpus=2) + ray.init(num_cpus=4) config = es.DEFAULT_CONFIG.copy() # Keep it simple. config["model"]["fcnet_hiddens"] = [10] @@ -18,6 +18,9 @@ class TestES(unittest.TestCase): config["num_workers"] = 1 config["episodes_per_batch"] = 10 config["train_batch_size"] = 100 + # Test eval workers ("normal" Trainer eval WorkerSet, unusual for ARS). + config["evaluation_interval"] = 1 + config["evaluation_num_workers"] = 2 num_iterations = 1 diff --git a/rllib/agents/pg/tests/test_pg.py b/rllib/agents/pg/tests/test_pg.py index 500e29cbd..55f0b2d70 100644 --- a/rllib/agents/pg/tests/test_pg.py +++ b/rllib/agents/pg/tests/test_pg.py @@ -21,7 +21,7 @@ class TestPG(unittest.TestCase): def test_pg_compilation(self): """Test whether a PGTrainer can be built with both frameworks.""" config = pg.DEFAULT_CONFIG.copy() - config["num_workers"] = 0 # Run locally. + config["num_workers"] = 0 num_iterations = 2 for _ in framework_iterator(config): diff --git a/rllib/agents/trainer.py b/rllib/agents/trainer.py index 1a48cb90f..8fad2cddb 100644 --- a/rllib/agents/trainer.py +++ b/rllib/agents/trainer.py @@ -14,6 +14,7 @@ from ray.exceptions import RayError from ray.rllib.agents.callbacks import DefaultCallbacks from ray.rllib.env.normalize_actions import NormalizeActionWrapper from ray.rllib.env.env_context import EnvContext +from ray.rllib.evaluation.rollout_worker import RolloutWorker from ray.rllib.models import MODEL_DEFAULTS from ray.rllib.policy import Policy from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID @@ -712,13 +713,10 @@ class Trainer(Trainable): Note that this default implementation does not do anything beyond merging evaluation_config with the normal trainer config. """ + # Call the `_before_evaluate` hook. self._before_evaluate() - - # Broadcast the new policy weights to all evaluation workers. - logger.info("Synchronizing weights to evaluation workers.") - weights = ray.put(self.workers.local_worker().save()) - self.evaluation_workers.foreach_worker( - lambda w: w.restore(ray.get(weights))) + # Sync weights to the evaluation WorkerSet. + self._sync_weights_to_workers(worker_set=self.evaluation_workers) self._sync_filters_if_needed(self.evaluation_workers) if self.config["custom_eval_function"]: @@ -759,6 +757,20 @@ class Trainer(Trainable): """Pre-evaluation callback.""" pass + @DeveloperAPI + def _sync_weights_to_workers( + self, + *, + worker_set: Optional[WorkerSet] = None, + workers: Optional[List[RolloutWorker]] = None, + ) -> None: + """Sync "main" weights to given WorkerSet or list of workers.""" + assert worker_set is not None + # Broadcast the new policy weights to all evaluation workers. + logger.info("Synchronizing weights to evaluation workers.") + weights = ray.put(self.workers.local_worker().save()) + worker_set.foreach_worker(lambda w: w.restore(ray.get(weights))) + @PublicAPI def compute_action(self, observation: TensorStructType, diff --git a/rllib/models/preprocessors.py b/rllib/models/preprocessors.py index c31cc533f..cb47daab1 100644 --- a/rllib/models/preprocessors.py +++ b/rllib/models/preprocessors.py @@ -58,6 +58,7 @@ class Preprocessor: observation = np.array(observation) try: if not self._obs_space.contains(observation): + print() raise ValueError( "Observation ({}) outside given space ({})!", observation, self._obs_space) diff --git a/rllib/policy/torch_policy_template.py b/rllib/policy/torch_policy_template.py index 1f98deb76..23601a6d3 100644 --- a/rllib/policy/torch_policy_template.py +++ b/rllib/policy/torch_policy_template.py @@ -22,9 +22,9 @@ torch, _ = try_import_torch() def build_torch_policy( name: str, *, - loss_fn: Callable[[ + loss_fn: Optional[Callable[[ Policy, ModelV2, Type[TorchDistributionWrapper], SampleBatch - ], Union[TensorType, List[TensorType]]], + ], Union[TensorType, List[TensorType]]]], get_default_config: Optional[Callable[[], TrainerConfigDict]] = None, stats_fn: Optional[Callable[[Policy, SampleBatch], Dict[ str, TensorType]]] = None, @@ -72,8 +72,9 @@ def build_torch_policy( Args: name (str): name of the policy (e.g., "PPOTorchPolicy") - loss_fn (Callable[[Policy, ModelV2, type, SampleBatch], TensorType]): - Callable that returns a loss tensor. + loss_fn (Optional[Callable[[Policy, ModelV2, + Type[TorchDistributionWrapper], SampleBatch], Union[TensorType, + List[TensorType]]]]): Callable that returns a loss tensor. get_default_config (Optional[Callable[[None], TrainerConfigDict]]): Optional callable that returns the default config to merge with any overrides. If None, uses only(!) the user-provided diff --git a/rllib/tests/test_evaluators.py b/rllib/tests/test_evaluators.py index 93399fa90..5b513fd1b 100644 --- a/rllib/tests/test_evaluators.py +++ b/rllib/tests/test_evaluators.py @@ -48,7 +48,7 @@ class EvalTest(unittest.TestCase): "framework": fw, }) # Given evaluation_interval=2, r0, r2, r4 should not contain - # evaluation metrics while r1, r3 should do. + # evaluation metrics, while r1, r3 should. r0 = agent.train() r1 = agent.train() r2 = agent.train()