[RLlib] ARS/ES eval workers not working: Issue 9933. (#11308)

This commit is contained in:
Sven Mika
2020-10-12 13:49:48 -07:00
committed by GitHub
parent 8ea1bc5ff9
commit 0c0f67c14d
12 changed files with 134 additions and 58 deletions
+23 -7
View File
@@ -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
+12 -12
View File
@@ -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()}
+5 -1
View File
@@ -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
+34 -8
View File
@@ -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):
+17 -14
View File
@@ -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):
+13 -3
View File
@@ -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):
+4 -1
View File
@@ -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
+1 -1
View File
@@ -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):
+18 -6
View File
@@ -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,
+1
View File
@@ -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)
+5 -4
View File
@@ -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
+1 -1
View File
@@ -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()