diff --git a/python/ray/rllib/agents/a3c/a2c.py b/python/ray/rllib/agents/a3c/a2c.py index f4e7f394a..c344592b9 100644 --- a/python/ray/rllib/agents/a3c/a2c.py +++ b/python/ray/rllib/agents/a3c/a2c.py @@ -4,6 +4,7 @@ from __future__ import print_function from ray.rllib.agents.a3c.a3c import A3CAgent, DEFAULT_CONFIG as A3C_CONFIG from ray.rllib.optimizers import SyncSamplesOptimizer +from ray.rllib.utils.annotations import override from ray.rllib.utils import merge_dicts A2C_DEFAULT_CONFIG = merge_dicts( @@ -22,6 +23,7 @@ class A2CAgent(A3CAgent): _agent_name = "A2C" _default_config = A2C_DEFAULT_CONFIG + @override(A3CAgent) def _make_optimizer(self): return SyncSamplesOptimizer(self.local_evaluator, self.remote_evaluators, diff --git a/python/ray/rllib/agents/a3c/a3c.py b/python/ray/rllib/agents/a3c/a3c.py index ebfec99e3..43daa0b3e 100644 --- a/python/ray/rllib/agents/a3c/a3c.py +++ b/python/ray/rllib/agents/a3c/a3c.py @@ -7,6 +7,7 @@ import time from ray.rllib.agents.a3c.a3c_tf_policy_graph import A3CPolicyGraph from ray.rllib.agents.agent import Agent, with_common_config from ray.rllib.optimizers import AsyncGradientsOptimizer +from ray.rllib.utils.annotations import override # yapf: disable # __sphinx_doc_begin__ @@ -44,6 +45,7 @@ class A3CAgent(Agent): _default_config = DEFAULT_CONFIG _policy_graph = A3CPolicyGraph + @override(Agent) def _init(self): if self.config["use_pytorch"]: from ray.rllib.agents.a3c.a3c_torch_policy_graph import \ @@ -58,11 +60,7 @@ class A3CAgent(Agent): self.env_creator, policy_cls, self.config["num_workers"]) self.optimizer = self._make_optimizer() - def _make_optimizer(self): - return AsyncGradientsOptimizer(self.local_evaluator, - self.remote_evaluators, - self.config["optimizer"]) - + @override(Agent) def _train(self): prev_steps = self.optimizer.num_steps_sampled start = time.time() @@ -73,3 +71,8 @@ class A3CAgent(Agent): result.update(timesteps_this_iter=self.optimizer.num_steps_sampled - prev_steps) return result + + def _make_optimizer(self): + return AsyncGradientsOptimizer(self.local_evaluator, + self.remote_evaluators, + self.config["optimizer"]) diff --git a/python/ray/rllib/agents/a3c/a3c_tf_policy_graph.py b/python/ray/rllib/agents/a3c/a3c_tf_policy_graph.py index 8aa60645a..50258f58a 100644 --- a/python/ray/rllib/agents/a3c/a3c_tf_policy_graph.py +++ b/python/ray/rllib/agents/a3c/a3c_tf_policy_graph.py @@ -10,10 +10,12 @@ import gym import ray from ray.rllib.utils.error import UnsupportedSpaceException from ray.rllib.utils.explained_variance import explained_variance +from ray.rllib.evaluation.policy_graph import PolicyGraph from ray.rllib.evaluation.postprocessing import compute_advantages from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph, \ LearningRateSchedule from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.utils.annotations import override class A3CLoss(object): @@ -118,30 +120,11 @@ class A3CPolicyGraph(LearningRateSchedule, TFPolicyGraph): self.sess.run(tf.global_variables_initializer()) - def extra_compute_action_fetches(self): - return {"vf_preds": self.vf} - - def value(self, ob, *args): - feed_dict = {self.observations: [ob], self.model.seq_lens: [1]} - assert len(args) == len(self.model.state_in), \ - (args, self.model.state_in) - for k, v in zip(self.model.state_in, args): - feed_dict[k] = v - vf = self.sess.run(self.vf, feed_dict) - return vf[0] - - def gradients(self, optimizer): - grads = tf.gradients(self.loss.total_loss, self.var_list) - self.grads, _ = tf.clip_by_global_norm(grads, self.config["grad_clip"]) - clipped_grads = list(zip(self.grads, self.var_list)) - return clipped_grads - - def extra_compute_grad_fetches(self): - return self.stats_fetches - + @override(PolicyGraph) def get_initial_state(self): return self.model.state_init + @override(PolicyGraph) def postprocess_trajectory(self, sample_batch, other_agent_batches=None, @@ -153,6 +136,30 @@ class A3CPolicyGraph(LearningRateSchedule, TFPolicyGraph): next_state = [] for i in range(len(self.model.state_in)): next_state.append([sample_batch["state_out_{}".format(i)][-1]]) - last_r = self.value(sample_batch["new_obs"][-1], *next_state) + last_r = self._value(sample_batch["new_obs"][-1], *next_state) return compute_advantages(sample_batch, last_r, self.config["gamma"], self.config["lambda"]) + + @override(TFPolicyGraph) + def gradients(self, optimizer): + grads = tf.gradients(self.loss.total_loss, self.var_list) + self.grads, _ = tf.clip_by_global_norm(grads, self.config["grad_clip"]) + clipped_grads = list(zip(self.grads, self.var_list)) + return clipped_grads + + @override(TFPolicyGraph) + def extra_compute_grad_fetches(self): + return self.stats_fetches + + @override(TFPolicyGraph) + def extra_compute_action_fetches(self): + return {"vf_preds": self.vf} + + def _value(self, ob, *args): + feed_dict = {self.observations: [ob], self.model.seq_lens: [1]} + assert len(args) == len(self.model.state_in), \ + (args, self.model.state_in) + for k, v in zip(self.model.state_in, args): + feed_dict[k] = v + vf = self.sess.run(self.vf, feed_dict) + return vf[0] diff --git a/python/ray/rllib/agents/a3c/a3c_torch_policy_graph.py b/python/ray/rllib/agents/a3c/a3c_torch_policy_graph.py index 3eecc3bb1..c24340d8d 100644 --- a/python/ray/rllib/agents/a3c/a3c_torch_policy_graph.py +++ b/python/ray/rllib/agents/a3c/a3c_torch_policy_graph.py @@ -10,7 +10,9 @@ import ray from ray.rllib.models.pytorch.misc import var_to_np from ray.rllib.models.catalog import ModelCatalog from ray.rllib.evaluation.postprocessing import compute_advantages +from ray.rllib.evaluation.policy_graph import PolicyGraph from ray.rllib.evaluation.torch_policy_graph import TorchPolicyGraph +from ray.rllib.utils.annotations import override class A3CLoss(nn.Module): @@ -56,12 +58,15 @@ class A3CTorchPolicyGraph(TorchPolicyGraph): loss, loss_inputs=["obs", "actions", "advantages", "value_targets"]) + @override(TorchPolicyGraph) def extra_action_out(self, model_out): return {"vf_preds": var_to_np(model_out[1])} + @override(TorchPolicyGraph) def optimizer(self): return torch.optim.Adam(self.model.parameters(), lr=self.config["lr"]) + @override(PolicyGraph) def postprocess_trajectory(self, sample_batch, other_agent_batches=None, diff --git a/python/ray/rllib/agents/agent.py b/python/ray/rllib/agents/agent.py index b84154f5b..8e6797eed 100644 --- a/python/ray/rllib/agents/agent.py +++ b/python/ray/rllib/agents/agent.py @@ -15,6 +15,7 @@ import ray from ray.rllib.models import MODEL_DEFAULTS from ray.rllib.evaluation.policy_evaluator import PolicyEvaluator from ray.rllib.optimizers.policy_optimizer import PolicyOptimizer +from ray.rllib.utils.annotations import override from ray.rllib.utils import FilterManager, deep_update, merge_dicts from ray.tune.registry import ENV_CREATOR, register_env, _global_registry from ray.tune.trainable import Trainable @@ -166,7 +167,48 @@ class Agent(Trainable): "tf_session_args", "env_config", "model", "optimizer", "multiagent" ] + def __init__(self, config=None, env=None, logger_creator=None): + """Initialize an RLLib agent. + + Args: + config (dict): Algorithm-specific configuration data. + env (str): Name of the environment to use. Note that this can also + be specified as the `env` key in config. + logger_creator (func): Function that creates a ray.tune.Logger + object. If unspecified, a default logger is created. + """ + + config = config or {} + Agent._validate_config(config) + + # Vars to synchronize to evaluators on each train call + self.global_vars = {"timestep": 0} + + # Agents allow env ids to be passed directly to the constructor. + self._env_id = _register_if_needed(env or config.get("env")) + + # Create a default logger creator if no logger_creator is specified + if logger_creator is None: + timestr = datetime.today().strftime("%Y-%m-%d_%H-%M-%S") + logdir_prefix = "{}_{}_{}".format(self._agent_name, self._env_id, + timestr) + + def default_logger_creator(config): + """Creates a Unified logger with a default logdir prefix + containing the agent name and the env id + """ + if not os.path.exists(DEFAULT_RESULTS_DIR): + os.makedirs(DEFAULT_RESULTS_DIR) + logdir = tempfile.mkdtemp( + prefix=logdir_prefix, dir=DEFAULT_RESULTS_DIR) + return UnifiedLogger(config, logdir, None) + + logger_creator = default_logger_creator + + Trainable.__init__(self, config, logger_creator) + @classmethod + @override(Trainable) def default_resource_request(cls, config): cf = dict(cls._default_config, **config) Agent._validate_config(cf) @@ -177,6 +219,147 @@ class Agent(Trainable): extra_cpu=cf["num_cpus_per_worker"] * cf["num_workers"], extra_gpu=cf["num_gpus_per_worker"] * cf["num_workers"]) + @override(Trainable) + def train(self): + """Overrides super.train to synchronize global vars.""" + + if hasattr(self, "optimizer") and isinstance(self.optimizer, + PolicyOptimizer): + self.global_vars["timestep"] = self.optimizer.num_steps_sampled + self.optimizer.local_evaluator.set_global_vars(self.global_vars) + for ev in self.optimizer.remote_evaluators: + ev.set_global_vars.remote(self.global_vars) + logger.debug("updated global vars: {}".format(self.global_vars)) + + if (self.config.get("observation_filter", "NoFilter") != "NoFilter" + and hasattr(self, "local_evaluator")): + FilterManager.synchronize( + self.local_evaluator.filters, + self.remote_evaluators, + update_remote=self.config["synchronize_filters"]) + logger.debug("synchronized filters: {}".format( + self.local_evaluator.filters)) + + result = Trainable.train(self) + if self.config["callbacks"].get("on_train_result"): + self.config["callbacks"]["on_train_result"]({ + "agent": self, + "result": result, + }) + return result + + @override(Trainable) + def _setup(self, config): + env = self._env_id + if env: + config["env"] = env + if _global_registry.contains(ENV_CREATOR, env): + self.env_creator = _global_registry.get(ENV_CREATOR, env) + else: + import gym # soft dependency + self.env_creator = lambda env_config: gym.make(env) + else: + self.env_creator = lambda env_config: None + + # Merge the supplied config with the class default + merged_config = copy.deepcopy(self._default_config) + merged_config = deep_update(merged_config, config, + self._allow_unknown_configs, + self._allow_unknown_subkeys) + self.config = merged_config + if self.config.get("log_level"): + logging.getLogger("ray.rllib").setLevel(self.config["log_level"]) + + # TODO(ekl) setting the graph is unnecessary for PyTorch agents + with tf.Graph().as_default(): + self._init() + + @override(Trainable) + def _stop(self): + # workaround for https://github.com/ray-project/ray/issues/1516 + if hasattr(self, "remote_evaluators"): + for ev in self.remote_evaluators: + ev.__ray_terminate__.remote() + if hasattr(self, "optimizer"): + self.optimizer.stop() + + @override(Trainable) + def _save(self, checkpoint_dir): + checkpoint_path = os.path.join(checkpoint_dir, + "checkpoint-{}".format(self.iteration)) + pickle.dump(self.__getstate__(), open(checkpoint_path, "wb")) + return checkpoint_path + + @override(Trainable) + def _restore(self, checkpoint_path): + extra_data = pickle.load(open(checkpoint_path, "rb")) + self.__setstate__(extra_data) + + def _init(self): + """Subclasses should override this for custom initialization.""" + + raise NotImplementedError + + def compute_action(self, observation, state=None, policy_id="default"): + """Computes an action for the specified policy. + + Arguments: + observation (obj): observation from the environment. + state (list): RNN hidden state, if any. If state is not None, + then all of compute_single_action(...) is returned + (computed action, rnn state, logits dictionary). + Otherwise compute_single_action(...)[0] is + returned (computed action). + policy_id (str): policy to query (only applies to multi-agent). + """ + + if state is None: + state = [] + filtered_obs = self.local_evaluator.filters[policy_id]( + observation, update=False) + if state: + return self.local_evaluator.for_policy( + lambda p: p.compute_single_action(filtered_obs, state), + policy_id=policy_id) + return self.local_evaluator.for_policy( + lambda p: p.compute_single_action(filtered_obs, state)[0], + policy_id=policy_id) + + @property + def iteration(self): + """Current training iter, auto-incremented with each train() call.""" + + return self._iteration + + @property + def _agent_name(self): + """Subclasses should override this to declare their name.""" + + raise NotImplementedError + + @property + def _default_config(self): + """Subclasses should override this to declare their default config.""" + + raise NotImplementedError + + def get_weights(self, policies=None): + """Return a dictionary of policy ids to weights. + + Arguments: + policies (list): Optional list of policies to return weights for, + or None for all policies. + """ + return self.local_evaluator.get_weights(policies) + + def set_weights(self, weights): + """Set policy weights by policy id. + + Arguments: + weights (dict): Map of policy ids to weights to set. + """ + self.local_evaluator.set_weights(weights) + def make_local_evaluator(self, env_creator, policy_graph): """Convenience method to return configured local evaluator.""" @@ -261,172 +444,6 @@ class Agent(Trainable): "The `use_gpu_for_workers` config is deprecated, please use " "`num_gpus_per_worker=1` instead.") - def __init__(self, config=None, env=None, logger_creator=None): - """Initialize an RLLib agent. - - Args: - config (dict): Algorithm-specific configuration data. - env (str): Name of the environment to use. Note that this can also - be specified as the `env` key in config. - logger_creator (func): Function that creates a ray.tune.Logger - object. If unspecified, a default logger is created. - """ - - config = config or {} - Agent._validate_config(config) - - # Vars to synchronize to evaluators on each train call - self.global_vars = {"timestep": 0} - - # Agents allow env ids to be passed directly to the constructor. - self._env_id = _register_if_needed(env or config.get("env")) - - # Create a default logger creator if no logger_creator is specified - if logger_creator is None: - timestr = datetime.today().strftime("%Y-%m-%d_%H-%M-%S") - logdir_prefix = "{}_{}_{}".format(self._agent_name, self._env_id, - timestr) - - def default_logger_creator(config): - """Creates a Unified logger with a default logdir prefix - containing the agent name and the env id - """ - if not os.path.exists(DEFAULT_RESULTS_DIR): - os.makedirs(DEFAULT_RESULTS_DIR) - logdir = tempfile.mkdtemp( - prefix=logdir_prefix, dir=DEFAULT_RESULTS_DIR) - return UnifiedLogger(config, logdir, None) - - logger_creator = default_logger_creator - - Trainable.__init__(self, config, logger_creator) - - def train(self): - """Overrides super.train to synchronize global vars.""" - - if hasattr(self, "optimizer") and isinstance(self.optimizer, - PolicyOptimizer): - self.global_vars["timestep"] = self.optimizer.num_steps_sampled - self.optimizer.local_evaluator.set_global_vars(self.global_vars) - for ev in self.optimizer.remote_evaluators: - ev.set_global_vars.remote(self.global_vars) - logger.debug("updated global vars: {}".format(self.global_vars)) - - if (self.config.get("observation_filter", "NoFilter") != "NoFilter" - and hasattr(self, "local_evaluator")): - FilterManager.synchronize( - self.local_evaluator.filters, - self.remote_evaluators, - update_remote=self.config["synchronize_filters"]) - logger.debug("synchronized filters: {}".format( - self.local_evaluator.filters)) - - result = Trainable.train(self) - if self.config["callbacks"].get("on_train_result"): - self.config["callbacks"]["on_train_result"]({ - "agent": self, - "result": result, - }) - return result - - def _setup(self, config): - env = self._env_id - if env: - config["env"] = env - if _global_registry.contains(ENV_CREATOR, env): - self.env_creator = _global_registry.get(ENV_CREATOR, env) - else: - import gym # soft dependency - self.env_creator = lambda env_config: gym.make(env) - else: - self.env_creator = lambda env_config: None - - # Merge the supplied config with the class default - merged_config = copy.deepcopy(self._default_config) - merged_config = deep_update(merged_config, config, - self._allow_unknown_configs, - self._allow_unknown_subkeys) - self.config = merged_config - if self.config.get("log_level"): - logging.getLogger("ray.rllib").setLevel(self.config["log_level"]) - - # TODO(ekl) setting the graph is unnecessary for PyTorch agents - with tf.Graph().as_default(): - self._init() - - def _init(self): - """Subclasses should override this for custom initialization.""" - - raise NotImplementedError - - @property - def iteration(self): - """Current training iter, auto-incremented with each train() call.""" - - return self._iteration - - @property - def _agent_name(self): - """Subclasses should override this to declare their name.""" - - raise NotImplementedError - - @property - def _default_config(self): - """Subclasses should override this to declare their default config.""" - - raise NotImplementedError - - def compute_action(self, observation, state=None, policy_id="default"): - """Computes an action for the specified policy. - - Arguments: - observation (obj): observation from the environment. - state (list): RNN hidden state, if any. If state is not None, - then all of compute_single_action(...) is returned - (computed action, rnn state, logits dictionary). - Otherwise compute_single_action(...)[0] is - returned (computed action). - policy_id (str): policy to query (only applies to multi-agent). - """ - - if state is None: - state = [] - filtered_obs = self.local_evaluator.filters[policy_id]( - observation, update=False) - if state: - return self.local_evaluator.for_policy( - lambda p: p.compute_single_action(filtered_obs, state), - policy_id=policy_id) - return self.local_evaluator.for_policy( - lambda p: p.compute_single_action(filtered_obs, state)[0], - policy_id=policy_id) - - def get_weights(self, policies=None): - """Return a dictionary of policy ids to weights. - - Arguments: - policies (list): Optional list of policies to return weights for, - or None for all policies. - """ - return self.local_evaluator.get_weights(policies) - - def set_weights(self, weights): - """Set policy weights by policy id. - - Arguments: - weights (dict): Map of policy ids to weights to set. - """ - self.local_evaluator.set_weights(weights) - - def _stop(self): - # workaround for https://github.com/ray-project/ray/issues/1516 - if hasattr(self, "remote_evaluators"): - for ev in self.remote_evaluators: - ev.__ray_terminate__.remote() - if hasattr(self, "optimizer"): - self.optimizer.stop() - def __getstate__(self): state = {} if hasattr(self, "local_evaluator"): @@ -444,16 +461,6 @@ class Agent(Trainable): if "optimizer" in state: self.optimizer.restore(state["optimizer"]) - def _save(self, checkpoint_dir): - checkpoint_path = os.path.join(checkpoint_dir, - "checkpoint-{}".format(self.iteration)) - pickle.dump(self.__getstate__(), open(checkpoint_path, "wb")) - return checkpoint_path - - def _restore(self, checkpoint_path): - extra_data = pickle.load(open(checkpoint_path, "rb")) - self.__setstate__(extra_data) - def _register_if_needed(env_object): if isinstance(env_object, six.string_types): diff --git a/python/ray/rllib/agents/ars/ars.py b/python/ray/rllib/agents/ars/ars.py index 67e9ba242..1b39a79d0 100644 --- a/python/ray/rllib/agents/ars/ars.py +++ b/python/ray/rllib/agents/ars/ars.py @@ -17,6 +17,7 @@ from ray.rllib.agents import Agent, with_common_config from ray.rllib.agents.ars import optimizers from ray.rllib.agents.ars import policies from ray.rllib.agents.ars import utils +from ray.rllib.utils.annotations import override from ray.rllib.utils import FilterManager logger = logging.getLogger(__name__) @@ -161,6 +162,7 @@ class ARSAgent(Agent): _agent_name = "ARS" _default_config = DEFAULT_CONFIG + @override(Agent) def _init(self): env = self.env_creator(self.config["env_config"]) from ray.rllib import models @@ -193,28 +195,7 @@ class ARSAgent(Agent): self.reward_list = [] self.tstart = time.time() - def _collect_results(self, theta_id, min_episodes): - num_episodes, num_timesteps = 0, 0 - results = [] - while num_episodes < min_episodes: - logger.info( - "Collected {} episodes {} timesteps so far this iter".format( - num_episodes, num_timesteps)) - rollout_ids = [ - worker.do_rollouts.remote(theta_id) for worker in self.workers - ] - # Get the results of the rollouts. - for result in ray.get(rollout_ids): - results.append(result) - # Update the number of episodes and the number of timesteps - # keeping in mind that result.noisy_lengths is a list of lists, - # where the inner lists have length 2. - num_episodes += sum(len(pair) for pair in result.noisy_lengths) - num_timesteps += sum( - sum(pair) for pair in result.noisy_lengths) - - return results, num_episodes, num_timesteps - + @override(Agent) def _train(self): config = self.config @@ -310,11 +291,38 @@ class ARSAgent(Agent): return result + @override(Agent) def _stop(self): # workaround for https://github.com/ray-project/ray/issues/1516 for w in self.workers: w.__ray_terminate__.remote() + @override(Agent) + def compute_action(self, observation): + return self.policy.compute(observation, update=True)[0] + + def _collect_results(self, theta_id, min_episodes): + num_episodes, num_timesteps = 0, 0 + results = [] + while num_episodes < min_episodes: + logger.info( + "Collected {} episodes {} timesteps so far this iter".format( + num_episodes, num_timesteps)) + rollout_ids = [ + worker.do_rollouts.remote(theta_id) for worker in self.workers + ] + # Get the results of the rollouts. + for result in ray.get(rollout_ids): + results.append(result) + # Update the number of episodes and the number of timesteps + # keeping in mind that result.noisy_lengths is a list of lists, + # where the inner lists have length 2. + num_episodes += sum(len(pair) for pair in result.noisy_lengths) + num_timesteps += sum( + sum(pair) for pair in result.noisy_lengths) + + return results, num_episodes, num_timesteps + def __getstate__(self): return { "weights": self.policy.get_weights(), @@ -329,6 +337,3 @@ class ARSAgent(Agent): FilterManager.synchronize({ "default": self.policy.get_filter() }, self.workers) - - def compute_action(self, observation): - return self.policy.compute(observation, update=True)[0] diff --git a/python/ray/rllib/agents/ddpg/apex.py b/python/ray/rllib/agents/ddpg/apex.py index c1699364a..6b3465013 100644 --- a/python/ray/rllib/agents/ddpg/apex.py +++ b/python/ray/rllib/agents/ddpg/apex.py @@ -3,6 +3,7 @@ from __future__ import division from __future__ import print_function from ray.rllib.agents.ddpg.ddpg import DDPGAgent, DEFAULT_CONFIG as DDPG_CONFIG +from ray.rllib.utils.annotations import override from ray.rllib.utils import merge_dicts APEX_DDPG_DEFAULT_CONFIG = merge_dicts( @@ -42,6 +43,7 @@ class ApexDDPGAgent(DDPGAgent): _agent_name = "APEX_DDPG" _default_config = APEX_DDPG_DEFAULT_CONFIG + @override(DDPGAgent) def update_target_if_needed(self): # Ape-X updates based on num steps trained, not sampled if self.optimizer.num_steps_trained - self.last_target_update_ts > \ diff --git a/python/ray/rllib/agents/ddpg/ddpg.py b/python/ray/rllib/agents/ddpg/ddpg.py index 564d8e12b..ca0e8087f 100644 --- a/python/ray/rllib/agents/ddpg/ddpg.py +++ b/python/ray/rllib/agents/ddpg/ddpg.py @@ -5,6 +5,7 @@ from __future__ import print_function from ray.rllib.agents.agent import with_common_config from ray.rllib.agents.dqn.dqn import DQNAgent from ray.rllib.agents.ddpg.ddpg_policy_graph import DDPGPolicyGraph +from ray.rllib.utils.annotations import override from ray.rllib.utils.schedules import ConstantSchedule, LinearSchedule OPTIMIZER_SHARED_CONFIGS = [ @@ -131,6 +132,7 @@ class DDPGAgent(DQNAgent): _default_config = DEFAULT_CONFIG _policy_graph = DDPGPolicyGraph + @override(DQNAgent) def _make_exploration_schedule(self, worker_index): # Override DQN's schedule to take into account `noise_scale` if self.config["per_worker_exploration"]: diff --git a/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py b/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py index eb5f14c2d..b8b625734 100644 --- a/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py +++ b/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py @@ -11,7 +11,9 @@ import ray from ray.rllib.agents.dqn.dqn_policy_graph import _huber_loss, \ _minimize_and_clip, _scope_vars, _postprocess_dqn from ray.rllib.models import ModelCatalog +from ray.rllib.utils.annotations import override from ray.rllib.utils.error import UnsupportedSpaceException +from ray.rllib.evaluation.policy_graph import PolicyGraph from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph A_SCOPE = "a_func" @@ -366,6 +368,75 @@ class DDPGPolicyGraph(TFPolicyGraph): # Hard initial update self.update_target(tau=1.0) + @override(TFPolicyGraph) + def optimizer(self): + return tf.train.AdamOptimizer(learning_rate=self.config["lr"]) + + @override(TFPolicyGraph) + def gradients(self, optimizer): + if self.config["grad_norm_clipping"] is not None: + actor_grads_and_vars = _minimize_and_clip( + optimizer, + self.loss.actor_loss, + var_list=self.p_func_vars, + clip_val=self.config["grad_norm_clipping"]) + critic_grads_and_vars = _minimize_and_clip( + optimizer, + self.loss.critic_loss, + var_list=self.q_func_vars + self.twin_q_func_vars + if self.config["twin_q"] else self.q_func_vars, + clip_val=self.config["grad_norm_clipping"]) + else: + actor_grads_and_vars = optimizer.compute_gradients( + self.loss.actor_loss, var_list=self.p_func_vars) + critic_grads_and_vars = optimizer.compute_gradients( + self.loss.critic_loss, + var_list=self.q_func_vars + self.twin_q_func_vars + if self.config["twin_q"] else self.q_func_vars) + actor_grads_and_vars = [(g, v) for (g, v) in actor_grads_and_vars + if g is not None] + critic_grads_and_vars = [(g, v) for (g, v) in critic_grads_and_vars + if g is not None] + grads_and_vars = actor_grads_and_vars + critic_grads_and_vars + return grads_and_vars + + @override(TFPolicyGraph) + def extra_compute_action_feed_dict(self): + return { + self.stochastic: True, + self.eps: self.cur_epsilon, + } + + @override(TFPolicyGraph) + def extra_compute_grad_fetches(self): + return { + "td_error": self.loss.td_error, + } + + @override(PolicyGraph) + def postprocess_trajectory(self, + sample_batch, + other_agent_batches=None, + episode=None): + return _postprocess_dqn(self, sample_batch) + + @override(TFPolicyGraph) + def get_weights(self): + return self.variables.get_weights() + + @override(TFPolicyGraph) + def set_weights(self, weights): + self.variables.set_weights(weights) + + @override(PolicyGraph) + def get_state(self): + return [TFPolicyGraph.get_state(self), self.cur_epsilon] + + @override(PolicyGraph) + def set_state(self, state): + TFPolicyGraph.set_state(self, state[0]) + self.set_epsilon(state[1]) + def _build_q_network(self, obs, obs_space, actions): q_net = QNetwork( ModelCatalog.get_model({ @@ -408,53 +479,6 @@ class DDPGPolicyGraph(TFPolicyGraph): self.config["use_huber"], self.config["huber_threshold"], self.config["twin_q"]) - def optimizer(self): - return tf.train.AdamOptimizer(learning_rate=self.config["lr"]) - - def gradients(self, optimizer): - if self.config["grad_norm_clipping"] is not None: - actor_grads_and_vars = _minimize_and_clip( - optimizer, - self.loss.actor_loss, - var_list=self.p_func_vars, - clip_val=self.config["grad_norm_clipping"]) - critic_grads_and_vars = _minimize_and_clip( - optimizer, - self.loss.critic_loss, - var_list=self.q_func_vars + self.twin_q_func_vars - if self.config["twin_q"] else self.q_func_vars, - clip_val=self.config["grad_norm_clipping"]) - else: - actor_grads_and_vars = optimizer.compute_gradients( - self.loss.actor_loss, var_list=self.p_func_vars) - critic_grads_and_vars = optimizer.compute_gradients( - self.loss.critic_loss, - var_list=self.q_func_vars + self.twin_q_func_vars - if self.config["twin_q"] else self.q_func_vars) - actor_grads_and_vars = [(g, v) for (g, v) in actor_grads_and_vars - if g is not None] - critic_grads_and_vars = [(g, v) for (g, v) in critic_grads_and_vars - if g is not None] - grads_and_vars = actor_grads_and_vars + critic_grads_and_vars - return grads_and_vars - - def extra_compute_action_feed_dict(self): - return { - self.stochastic: True, - self.eps: self.cur_epsilon, - } - - def extra_compute_grad_fetches(self): - return { - "td_error": self.loss.td_error, - } - - def postprocess_trajectory(self, - sample_batch, - other_agent_batches=None, - episode=None): - return _postprocess_dqn(self, sample_batch) - def compute_td_error(self, obs_t, act_t, rew_t, obs_tp1, done_mask, importance_weights): td_err = self.sess.run( @@ -480,16 +504,3 @@ class DDPGPolicyGraph(TFPolicyGraph): def set_epsilon(self, epsilon): self.cur_epsilon = epsilon - - def get_weights(self): - return self.variables.get_weights() - - def set_weights(self, weights): - self.variables.set_weights(weights) - - def get_state(self): - return [TFPolicyGraph.get_state(self), self.cur_epsilon] - - def set_state(self, state): - TFPolicyGraph.set_state(self, state[0]) - self.set_epsilon(state[1]) diff --git a/python/ray/rllib/agents/dqn/apex.py b/python/ray/rllib/agents/dqn/apex.py index a6738d661..c9b15e0ec 100644 --- a/python/ray/rllib/agents/dqn/apex.py +++ b/python/ray/rllib/agents/dqn/apex.py @@ -4,6 +4,7 @@ from __future__ import print_function from ray.rllib.agents.dqn.dqn import DQNAgent, DEFAULT_CONFIG as DQN_CONFIG from ray.rllib.utils import merge_dicts +from ray.rllib.utils.annotations import override # yapf: disable # __sphinx_doc_begin__ @@ -45,6 +46,7 @@ class ApexAgent(DQNAgent): _agent_name = "APEX" _default_config = APEX_DEFAULT_CONFIG + @override(DQNAgent) def update_target_if_needed(self): # Ape-X updates based on num steps trained, not sampled if self.optimizer.num_steps_trained - self.last_target_update_ts > \ diff --git a/python/ray/rllib/agents/dqn/dqn.py b/python/ray/rllib/agents/dqn/dqn.py index db5d261ee..10e3edd48 100644 --- a/python/ray/rllib/agents/dqn/dqn.py +++ b/python/ray/rllib/agents/dqn/dqn.py @@ -7,6 +7,7 @@ import time from ray.rllib import optimizers from ray.rllib.agents.agent import Agent, with_common_config from ray.rllib.agents.dqn.dqn_policy_graph import DQNPolicyGraph +from ray.rllib.utils.annotations import override from ray.rllib.utils.schedules import ConstantSchedule, LinearSchedule OPTIMIZER_SHARED_CONFIGS = [ @@ -117,6 +118,7 @@ class DQNAgent(Agent): _default_config = DEFAULT_CONFIG _policy_graph = DQNPolicyGraph + @override(Agent) def _init(self): # Update effective batch size to include n-step adjusted_batch_size = max(self.config["sample_batch_size"], @@ -159,43 +161,12 @@ class DQNAgent(Agent): # Create the remote evaluators *after* the replay actors if self.remote_evaluators is None: self.remote_evaluators = create_remote_evaluators() - self.optimizer.set_evaluators(self.remote_evaluators) + self.optimizer._set_evaluators(self.remote_evaluators) self.last_target_update_ts = 0 self.num_target_updates = 0 - def _make_exploration_schedule(self, worker_index): - # Use either a different `eps` per worker, or a linear schedule. - if self.config["per_worker_exploration"]: - assert self.config["num_workers"] > 1, \ - "This requires multiple workers" - if worker_index >= 0: - exponent = ( - 1 + - worker_index / float(self.config["num_workers"] - 1) * 7) - return ConstantSchedule(0.4**exponent) - else: - # local ev should have zero exploration so that eval rollouts - # run properly - return ConstantSchedule(0.0) - return LinearSchedule( - schedule_timesteps=int(self.config["exploration_fraction"] * - self.config["schedule_max_timesteps"]), - initial_p=1.0, - final_p=self.config["exploration_final_eps"]) - - @property - def global_timestep(self): - return self.optimizer.num_steps_sampled - - def update_target_if_needed(self): - if self.global_timestep - self.last_target_update_ts > \ - self.config["target_network_update_freq"]: - self.local_evaluator.foreach_trainable_policy( - lambda p, _: p.update_target()) - self.last_target_update_ts = self.global_timestep - self.num_target_updates += 1 - + @override(Agent) def _train(self): start_timestep = self.global_timestep @@ -236,6 +207,38 @@ class DQNAgent(Agent): }, **self.optimizer.stats())) return result + def update_target_if_needed(self): + if self.global_timestep - self.last_target_update_ts > \ + self.config["target_network_update_freq"]: + self.local_evaluator.foreach_trainable_policy( + lambda p, _: p.update_target()) + self.last_target_update_ts = self.global_timestep + self.num_target_updates += 1 + + @property + def global_timestep(self): + return self.optimizer.num_steps_sampled + + def _make_exploration_schedule(self, worker_index): + # Use either a different `eps` per worker, or a linear schedule. + if self.config["per_worker_exploration"]: + assert self.config["num_workers"] > 1, \ + "This requires multiple workers" + if worker_index >= 0: + exponent = ( + 1 + + worker_index / float(self.config["num_workers"] - 1) * 7) + return ConstantSchedule(0.4**exponent) + else: + # local ev should have zero exploration so that eval rollouts + # run properly + return ConstantSchedule(0.0) + return LinearSchedule( + schedule_timesteps=int(self.config["exploration_fraction"] * + self.config["schedule_max_timesteps"]), + initial_p=1.0, + final_p=self.config["exploration_final_eps"]) + def __getstate__(self): state = Agent.__getstate__(self) state.update({ diff --git a/python/ray/rllib/agents/dqn/dqn_policy_graph.py b/python/ray/rllib/agents/dqn/dqn_policy_graph.py index af43b7d17..625e577ff 100644 --- a/python/ray/rllib/agents/dqn/dqn_policy_graph.py +++ b/python/ray/rllib/agents/dqn/dqn_policy_graph.py @@ -10,7 +10,9 @@ import tensorflow.contrib.layers as layers import ray from ray.rllib.models import ModelCatalog from ray.rllib.evaluation.sample_batch import SampleBatch +from ray.rllib.utils.annotations import override from ray.rllib.utils.error import UnsupportedSpaceException +from ray.rllib.evaluation.policy_graph import PolicyGraph from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph Q_SCOPE = "q_func" @@ -390,6 +392,76 @@ class DQNPolicyGraph(TFPolicyGraph): update_ops=q_batchnorm_update_ops) self.sess.run(tf.global_variables_initializer()) + @override(TFPolicyGraph) + def optimizer(self): + return tf.train.AdamOptimizer( + learning_rate=self.config["lr"], + epsilon=self.config["adam_epsilon"]) + + @override(TFPolicyGraph) + def gradients(self, optimizer): + if self.config["grad_norm_clipping"] is not None: + grads_and_vars = _minimize_and_clip( + optimizer, + self.loss.loss, + var_list=self.q_func_vars, + clip_val=self.config["grad_norm_clipping"]) + else: + grads_and_vars = optimizer.compute_gradients( + self.loss.loss, var_list=self.q_func_vars) + grads_and_vars = [(g, v) for (g, v) in grads_and_vars if g is not None] + return grads_and_vars + + @override(TFPolicyGraph) + def extra_compute_action_feed_dict(self): + return { + self.stochastic: True, + self.eps: self.cur_epsilon, + } + + @override(TFPolicyGraph) + def extra_compute_grad_fetches(self): + return { + "td_error": self.loss.td_error, + "stats": self.loss.stats, + } + + @override(PolicyGraph) + def postprocess_trajectory(self, + sample_batch, + other_agent_batches=None, + episode=None): + return _postprocess_dqn(self, sample_batch) + + @override(PolicyGraph) + def get_state(self): + return [TFPolicyGraph.get_state(self), self.cur_epsilon] + + @override(PolicyGraph) + def set_state(self, state): + TFPolicyGraph.set_state(self, state[0]) + self.set_epsilon(state[1]) + + def compute_td_error(self, obs_t, act_t, rew_t, obs_tp1, done_mask, + importance_weights): + td_err = self.sess.run( + self.loss.td_error, + feed_dict={ + self.obs_t: [np.array(ob) for ob in obs_t], + self.act_t: act_t, + self.rew_t: rew_t, + self.obs_tp1: [np.array(ob) for ob in obs_tp1], + self.done_mask: done_mask, + self.importance_weights: importance_weights + }) + return td_err + + def update_target(self): + return self.sess.run(self.update_target_expr) + + def set_epsilon(self, epsilon): + self.cur_epsilon = epsilon + def _build_q_network(self, obs, space): qnet = QNetwork( ModelCatalog.get_model({ @@ -413,71 +485,8 @@ class DQNPolicyGraph(TFPolicyGraph): self.config["n_step"], self.config["num_atoms"], self.config["v_min"], self.config["v_max"]) - def optimizer(self): - return tf.train.AdamOptimizer( - learning_rate=self.config["lr"], - epsilon=self.config["adam_epsilon"]) - def gradients(self, optimizer): - if self.config["grad_norm_clipping"] is not None: - grads_and_vars = _minimize_and_clip( - optimizer, - self.loss.loss, - var_list=self.q_func_vars, - clip_val=self.config["grad_norm_clipping"]) - else: - grads_and_vars = optimizer.compute_gradients( - self.loss.loss, var_list=self.q_func_vars) - grads_and_vars = [(g, v) for (g, v) in grads_and_vars if g is not None] - return grads_and_vars - - def extra_compute_action_feed_dict(self): - return { - self.stochastic: True, - self.eps: self.cur_epsilon, - } - - def extra_compute_grad_fetches(self): - return { - "td_error": self.loss.td_error, - "stats": self.loss.stats, - } - - def postprocess_trajectory(self, - sample_batch, - other_agent_batches=None, - episode=None): - return _postprocess_dqn(self, sample_batch) - - def compute_td_error(self, obs_t, act_t, rew_t, obs_tp1, done_mask, - importance_weights): - td_err = self.sess.run( - self.loss.td_error, - feed_dict={ - self.obs_t: [np.array(ob) for ob in obs_t], - self.act_t: act_t, - self.rew_t: rew_t, - self.obs_tp1: [np.array(ob) for ob in obs_tp1], - self.done_mask: done_mask, - self.importance_weights: importance_weights - }) - return td_err - - def update_target(self): - return self.sess.run(self.update_target_expr) - - def set_epsilon(self, epsilon): - self.cur_epsilon = epsilon - - def get_state(self): - return [TFPolicyGraph.get_state(self), self.cur_epsilon] - - def set_state(self, state): - TFPolicyGraph.set_state(self, state[0]) - self.set_epsilon(state[1]) - - -def adjust_nstep(n_step, gamma, obs, actions, rewards, new_obs, dones): +def _adjust_nstep(n_step, gamma, obs, actions, rewards, new_obs, dones): """Rewrites the given trajectory fragments to encode n-step rewards. reward[i] = ( @@ -510,9 +519,9 @@ def _postprocess_dqn(policy_graph, sample_batch): # N-step Q adjustments if policy_graph.config["n_step"] > 1: - adjust_nstep(policy_graph.config["n_step"], - policy_graph.config["gamma"], obs, actions, rewards, - new_obs, dones) + _adjust_nstep(policy_graph.config["n_step"], + policy_graph.config["gamma"], obs, actions, rewards, + new_obs, dones) batch = SampleBatch({ "obs": obs, diff --git a/python/ray/rllib/agents/es/es.py b/python/ray/rllib/agents/es/es.py index 550296812..4aa4a86aa 100644 --- a/python/ray/rllib/agents/es/es.py +++ b/python/ray/rllib/agents/es/es.py @@ -16,6 +16,7 @@ from ray.rllib.agents import Agent, with_common_config from ray.rllib.agents.es import optimizers from ray.rllib.agents.es import policies from ray.rllib.agents.es import utils +from ray.rllib.utils.annotations import override from ray.rllib.utils import FilterManager logger = logging.getLogger(__name__) @@ -167,6 +168,7 @@ class ESAgent(Agent): _agent_name = "ES" _default_config = DEFAULT_CONFIG + @override(Agent) def _init(self): policy_params = {"action_noise_std": 0.01} @@ -198,28 +200,7 @@ class ESAgent(Agent): self.reward_list = [] self.tstart = time.time() - def _collect_results(self, theta_id, min_episodes, min_timesteps): - num_episodes, num_timesteps = 0, 0 - results = [] - while num_episodes < min_episodes or num_timesteps < min_timesteps: - logger.info( - "Collected {} episodes {} timesteps so far this iter".format( - num_episodes, num_timesteps)) - rollout_ids = [ - worker.do_rollouts.remote(theta_id) for worker in self.workers - ] - # Get the results of the rollouts. - for result in ray.get(rollout_ids): - results.append(result) - # Update the number of episodes and the number of timesteps - # keeping in mind that result.noisy_lengths is a list of lists, - # where the inner lists have length 2. - num_episodes += sum(len(pair) for pair in result.noisy_lengths) - num_timesteps += sum( - sum(pair) for pair in result.noisy_lengths) - - return results, num_episodes, num_timesteps - + @override(Agent) def _train(self): config = self.config @@ -307,11 +288,38 @@ class ESAgent(Agent): return result + @override(Agent) + def compute_action(self, observation): + return self.policy.compute(observation, update=False)[0] + + @override(Agent) def _stop(self): # workaround for https://github.com/ray-project/ray/issues/1516 for w in self.workers: w.__ray_terminate__.remote() + def _collect_results(self, theta_id, min_episodes, min_timesteps): + num_episodes, num_timesteps = 0, 0 + results = [] + while num_episodes < min_episodes or num_timesteps < min_timesteps: + logger.info( + "Collected {} episodes {} timesteps so far this iter".format( + num_episodes, num_timesteps)) + rollout_ids = [ + worker.do_rollouts.remote(theta_id) for worker in self.workers + ] + # Get the results of the rollouts. + for result in ray.get(rollout_ids): + results.append(result) + # Update the number of episodes and the number of timesteps + # keeping in mind that result.noisy_lengths is a list of lists, + # where the inner lists have length 2. + num_episodes += sum(len(pair) for pair in result.noisy_lengths) + num_timesteps += sum( + sum(pair) for pair in result.noisy_lengths) + + return results, num_episodes, num_timesteps + def __getstate__(self): return { "weights": self.policy.get_weights(), @@ -326,6 +334,3 @@ class ESAgent(Agent): FilterManager.synchronize({ "default": self.policy.get_filter() }, self.workers) - - def compute_action(self, observation): - return self.policy.compute(observation, update=False)[0] diff --git a/python/ray/rllib/agents/impala/impala.py b/python/ray/rllib/agents/impala/impala.py index 45af92200..aa789387f 100644 --- a/python/ray/rllib/agents/impala/impala.py +++ b/python/ray/rllib/agents/impala/impala.py @@ -8,6 +8,7 @@ from ray.rllib.agents.a3c.a3c_tf_policy_graph import A3CPolicyGraph from ray.rllib.agents.impala.vtrace_policy_graph import VTracePolicyGraph from ray.rllib.agents.agent import Agent, with_common_config from ray.rllib.optimizers import AsyncSamplesOptimizer +from ray.rllib.utils.annotations import override OPTIMIZER_SHARED_CONFIGS = [ "lr", @@ -77,6 +78,7 @@ class ImpalaAgent(Agent): _default_config = DEFAULT_CONFIG _policy_graph = VTracePolicyGraph + @override(Agent) def _init(self): for k in OPTIMIZER_SHARED_CONFIGS: if k not in self.config["optimizer"]: @@ -93,6 +95,7 @@ class ImpalaAgent(Agent): self.remote_evaluators, self.config["optimizer"]) + @override(Agent) def _train(self): prev_steps = self.optimizer.num_steps_sampled start = time.time() diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py index cfa2f1373..5eed0a6e7 100644 --- a/python/ray/rllib/agents/impala/vtrace_policy_graph.py +++ b/python/ray/rllib/agents/impala/vtrace_policy_graph.py @@ -11,9 +11,11 @@ import gym import ray from ray.rllib.agents.impala import vtrace +from ray.rllib.evaluation.policy_graph import PolicyGraph from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph, \ LearningRateSchedule from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.utils.annotations import override from ray.rllib.utils.error import UnsupportedSpaceException from ray.rllib.utils.explained_variance import explained_variance from ray.rllib.models.action_dist import Categorical @@ -242,6 +244,15 @@ class VTracePolicyGraph(LearningRateSchedule, TFPolicyGraph): }, } + @override(TFPolicyGraph) + def copy(self, existing_inputs): + return VTracePolicyGraph( + self.observation_space, + self.action_space, + self.config, + existing_inputs=existing_inputs) + + @override(TFPolicyGraph) def optimizer(self): if self.config["opt_type"] == "adam": return tf.train.AdamOptimizer(self.cur_lr) @@ -250,18 +261,22 @@ class VTracePolicyGraph(LearningRateSchedule, TFPolicyGraph): self.config["momentum"], self.config["epsilon"]) + @override(TFPolicyGraph) def gradients(self, optimizer): grads = tf.gradients(self.loss.total_loss, self.var_list) self.grads, _ = tf.clip_by_global_norm(grads, self.config["grad_clip"]) clipped_grads = list(zip(self.grads, self.var_list)) return clipped_grads + @override(TFPolicyGraph) def extra_compute_action_fetches(self): return {"behaviour_logits": self.model.outputs} + @override(TFPolicyGraph) def extra_compute_grad_fetches(self): return self.stats_fetches + @override(PolicyGraph) def postprocess_trajectory(self, sample_batch, other_agent_batches=None, @@ -269,12 +284,6 @@ class VTracePolicyGraph(LearningRateSchedule, TFPolicyGraph): del sample_batch.data["new_obs"] # not used, so save some bandwidth return sample_batch + @override(PolicyGraph) def get_initial_state(self): return self.model.state_init - - def copy(self, existing_inputs): - return VTracePolicyGraph( - self.observation_space, - self.action_space, - self.config, - existing_inputs=existing_inputs) diff --git a/python/ray/rllib/agents/pg/pg.py b/python/ray/rllib/agents/pg/pg.py index 925cbc1a1..69c676186 100644 --- a/python/ray/rllib/agents/pg/pg.py +++ b/python/ray/rllib/agents/pg/pg.py @@ -5,6 +5,7 @@ from __future__ import print_function from ray.rllib.agents.agent import Agent, with_common_config from ray.rllib.agents.pg.pg_policy_graph import PGPolicyGraph from ray.rllib.optimizers import SyncSamplesOptimizer +from ray.rllib.utils.annotations import override # yapf: disable # __sphinx_doc_begin__ @@ -29,6 +30,7 @@ class PGAgent(Agent): _default_config = DEFAULT_CONFIG _policy_graph = PGPolicyGraph + @override(Agent) def _init(self): self.local_evaluator = self.make_local_evaluator( self.env_creator, self._policy_graph) @@ -38,6 +40,7 @@ class PGAgent(Agent): self.remote_evaluators, self.config["optimizer"]) + @override(Agent) def _train(self): prev_steps = self.optimizer.num_steps_sampled self.optimizer.step() diff --git a/python/ray/rllib/agents/pg/pg_policy_graph.py b/python/ray/rllib/agents/pg/pg_policy_graph.py index 2a342c117..59e9a9eff 100644 --- a/python/ray/rllib/agents/pg/pg_policy_graph.py +++ b/python/ray/rllib/agents/pg/pg_policy_graph.py @@ -7,7 +7,9 @@ import tensorflow as tf import ray from ray.rllib.models.catalog import ModelCatalog from ray.rllib.evaluation.postprocessing import compute_advantages +from ray.rllib.evaluation.policy_graph import PolicyGraph from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph +from ray.rllib.utils.annotations import override class PGLoss(object): @@ -75,6 +77,7 @@ class PGPolicyGraph(TFPolicyGraph): max_seq_len=config["model"]["max_seq_len"]) sess.run(tf.global_variables_initializer()) + @override(PolicyGraph) def postprocess_trajectory(self, sample_batch, other_agent_batches=None, @@ -83,5 +86,6 @@ class PGPolicyGraph(TFPolicyGraph): return compute_advantages( sample_batch, 0.0, self.config["gamma"], use_gae=False) + @override(PolicyGraph) def get_initial_state(self): return self.model.state_init diff --git a/python/ray/rllib/agents/ppo/ppo.py b/python/ray/rllib/agents/ppo/ppo.py index d5e50832f..ec8fb9c0a 100644 --- a/python/ray/rllib/agents/ppo/ppo.py +++ b/python/ray/rllib/agents/ppo/ppo.py @@ -7,6 +7,7 @@ import logging from ray.rllib.agents import Agent, with_common_config from ray.rllib.agents.ppo.ppo_policy_graph import PPOPolicyGraph from ray.rllib.optimizers import SyncSamplesOptimizer, LocalMultiGPUOptimizer +from ray.rllib.utils.annotations import override logger = logging.getLogger(__name__) @@ -64,6 +65,7 @@ class PPOAgent(Agent): _default_config = DEFAULT_CONFIG _policy_graph = PPOPolicyGraph + @override(Agent) def _init(self): self._validate_config() self.local_evaluator = self.make_local_evaluator( @@ -86,6 +88,25 @@ class PPOAgent(Agent): "standardize_fields": ["advantages"], }) + @override(Agent) + def _train(self): + prev_steps = self.optimizer.num_steps_sampled + fetches = self.optimizer.step() + if "kl" in fetches: + # single-agent + self.local_evaluator.for_policy( + lambda pi: pi.update_kl(fetches["kl"])) + else: + # multi-agent + self.local_evaluator.foreach_trainable_policy( + lambda pi, pi_id: pi.update_kl(fetches[pi_id]["kl"])) + res = self.optimizer.collect_metrics( + self.config["collect_metrics_timeout"]) + res.update( + timesteps_this_iter=self.optimizer.num_steps_sampled - prev_steps, + info=dict(fetches, **res.get("info", {}))) + return res + def _validate_config(self): waste_ratio = ( self.config["sample_batch_size"] * self.config["num_workers"] / @@ -116,21 +137,3 @@ class PPOAgent(Agent): logger.warn( "By default, observations will be normalized with {}".format( self.config["observation_filter"])) - - def _train(self): - prev_steps = self.optimizer.num_steps_sampled - fetches = self.optimizer.step() - if "kl" in fetches: - # single-agent - self.local_evaluator.for_policy( - lambda pi: pi.update_kl(fetches["kl"])) - else: - # multi-agent - self.local_evaluator.foreach_trainable_policy( - lambda pi, pi_id: pi.update_kl(fetches[pi_id]["kl"])) - res = self.optimizer.collect_metrics( - self.config["collect_metrics_timeout"]) - res.update( - timesteps_this_iter=self.optimizer.num_steps_sampled - prev_steps, - info=dict(fetches, **res.get("info", {}))) - return res diff --git a/python/ray/rllib/agents/ppo/ppo_policy_graph.py b/python/ray/rllib/agents/ppo/ppo_policy_graph.py index 3762f16f9..852047976 100644 --- a/python/ray/rllib/agents/ppo/ppo_policy_graph.py +++ b/python/ray/rllib/agents/ppo/ppo_policy_graph.py @@ -6,9 +6,11 @@ import tensorflow as tf import ray from ray.rllib.evaluation.postprocessing import compute_advantages +from ray.rllib.evaluation.policy_graph import PolicyGraph from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph, \ LearningRateSchedule from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.utils.annotations import override from ray.rllib.utils.explained_variance import explained_variance @@ -254,6 +256,7 @@ class PPOPolicyGraph(LearningRateSchedule, TFPolicyGraph): "entropy": self.loss_obj.mean_entropy } + @override(TFPolicyGraph) def copy(self, existing_inputs): """Creates a copy of self using existing input placeholders.""" return PPOPolicyGraph( @@ -262,29 +265,7 @@ class PPOPolicyGraph(LearningRateSchedule, TFPolicyGraph): self.config, existing_inputs=existing_inputs) - def extra_compute_action_fetches(self): - return {"vf_preds": self.value_function, "logits": self.logits} - - def extra_compute_grad_fetches(self): - return self.stats_fetches - - def update_kl(self, sampled_kl): - if sampled_kl > 2.0 * self.kl_target: - self.kl_coeff_val *= 1.5 - elif sampled_kl < 0.5 * self.kl_target: - self.kl_coeff_val *= 0.5 - self.kl_coeff.load(self.kl_coeff_val, session=self.sess) - return self.kl_coeff_val - - def value(self, ob, *args): - feed_dict = {self.observations: [ob], self.model.seq_lens: [1]} - assert len(args) == len(self.model.state_in), \ - (args, self.model.state_in) - for k, v in zip(self.model.state_in, args): - feed_dict[k] = v - vf = self.sess.run(self.value_function, feed_dict) - return vf[0] - + @override(PolicyGraph) def postprocess_trajectory(self, sample_batch, other_agent_batches=None, @@ -296,7 +277,7 @@ class PPOPolicyGraph(LearningRateSchedule, TFPolicyGraph): next_state = [] for i in range(len(self.model.state_in)): next_state.append([sample_batch["state_out_{}".format(i)][-1]]) - last_r = self.value(sample_batch["new_obs"][-1], *next_state) + last_r = self._value(sample_batch["new_obs"][-1], *next_state) batch = compute_advantages( sample_batch, last_r, @@ -305,9 +286,36 @@ class PPOPolicyGraph(LearningRateSchedule, TFPolicyGraph): use_gae=self.config["use_gae"]) return batch + @override(TFPolicyGraph) def gradients(self, optimizer): return optimizer.compute_gradients( self._loss, colocate_gradients_with_ops=True) + @override(PolicyGraph) def get_initial_state(self): return self.model.state_init + + @override(TFPolicyGraph) + def extra_compute_action_fetches(self): + return {"vf_preds": self.value_function, "logits": self.logits} + + @override(TFPolicyGraph) + def extra_compute_grad_fetches(self): + return self.stats_fetches + + def update_kl(self, sampled_kl): + if sampled_kl > 2.0 * self.kl_target: + self.kl_coeff_val *= 1.5 + elif sampled_kl < 0.5 * self.kl_target: + self.kl_coeff_val *= 0.5 + self.kl_coeff.load(self.kl_coeff_val, session=self.sess) + return self.kl_coeff_val + + def _value(self, ob, *args): + feed_dict = {self.observations: [ob], self.model.seq_lens: [1]} + assert len(args) == len(self.model.state_in), \ + (args, self.model.state_in) + for k, v in zip(self.model.state_in, args): + feed_dict[k] = v + vf = self.sess.run(self.value_function, feed_dict) + return vf[0] diff --git a/python/ray/rllib/env/async_vector_env.py b/python/ray/rllib/env/async_vector_env.py index edbf7a233..72cd812de 100644 --- a/python/ray/rllib/env/async_vector_env.py +++ b/python/ray/rllib/env/async_vector_env.py @@ -5,6 +5,7 @@ from __future__ import print_function from ray.rllib.env.external_env import ExternalEnv from ray.rllib.env.vector_env import VectorEnv from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.utils.annotations import override class AsyncVectorEnv(object): @@ -158,6 +159,7 @@ class _ExternalEnvToAsync(AsyncVectorEnv): self.observation_space = external_env.observation_space external_env.start() + @override(AsyncVectorEnv) def poll(self): with self.external_env._results_avail_condition: results = self._poll() @@ -172,6 +174,12 @@ class _ExternalEnvToAsync(AsyncVectorEnv): "ExternalEnv was created with max_concurrent={}".format(limit)) return results + @override(AsyncVectorEnv) + def send_actions(self, action_dict): + for eid, action in action_dict.items(): + self.external_env._episodes[eid].action_queue.put( + action[_DUMMY_AGENT_ID]) + def _poll(self): all_obs, all_rewards, all_dones, all_infos = {}, {}, {}, {} off_policy_actions = {} @@ -195,11 +203,6 @@ class _ExternalEnvToAsync(AsyncVectorEnv): _with_dummy_agent_id(all_infos), \ _with_dummy_agent_id(off_policy_actions) - def send_actions(self, action_dict): - for eid, action in action_dict.items(): - self.external_env._episodes[eid].action_queue.put( - action[_DUMMY_AGENT_ID]) - class _VectorEnvToAsync(AsyncVectorEnv): """Internal adapter of VectorEnv to AsyncVectorEnv. @@ -219,6 +222,7 @@ class _VectorEnvToAsync(AsyncVectorEnv): self.cur_dones = [False for _ in range(self.num_envs)] self.cur_infos = [None for _ in range(self.num_envs)] + @override(AsyncVectorEnv) def poll(self): if self.new_obs is None: self.new_obs = self.vector_env.vector_reset() @@ -235,6 +239,7 @@ class _VectorEnvToAsync(AsyncVectorEnv): _with_dummy_agent_id(dones, "__all__"), \ _with_dummy_agent_id(infos), {} + @override(AsyncVectorEnv) def send_actions(self, action_dict): action_vector = [None] * self.num_envs for i in range(self.num_envs): @@ -242,9 +247,11 @@ class _VectorEnvToAsync(AsyncVectorEnv): self.new_obs, self.cur_rewards, self.cur_dones, self.cur_infos = \ self.vector_env.vector_step(action_vector) + @override(AsyncVectorEnv) def try_reset(self, env_id): return {_DUMMY_AGENT_ID: self.vector_env.reset_at(env_id)} + @override(AsyncVectorEnv) def get_unwrapped(self): return self.vector_env.get_unwrapped() @@ -275,12 +282,14 @@ class _MultiAgentEnvToAsync(AsyncVectorEnv): assert isinstance(env, MultiAgentEnv) self.env_states = [_MultiAgentEnvState(env) for env in self.envs] + @override(AsyncVectorEnv) def poll(self): obs, rewards, dones, infos = {}, {}, {}, {} for i, env_state in enumerate(self.env_states): obs[i], rewards[i], dones[i], infos[i] = env_state.poll() return obs, rewards, dones, infos, {} + @override(AsyncVectorEnv) def send_actions(self, action_dict): for env_id, agent_dict in action_dict.items(): if env_id in self.dones: @@ -302,6 +311,7 @@ class _MultiAgentEnvToAsync(AsyncVectorEnv): self.dones.add(env_id) self.env_states[env_id].observe(obs, rewards, dones, infos) + @override(AsyncVectorEnv) def try_reset(self, env_id): obs = self.env_states[env_id].reset() assert isinstance(obs, dict), "Not a multi-agent obs" diff --git a/python/ray/rllib/env/vector_env.py b/python/ray/rllib/env/vector_env.py index 8d2289cf4..c2eb16920 100644 --- a/python/ray/rllib/env/vector_env.py +++ b/python/ray/rllib/env/vector_env.py @@ -2,6 +2,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from ray.rllib.utils.annotations import override + class VectorEnv(object): """An environment that supports batch evaluation. @@ -72,12 +74,15 @@ class _VectorizedGymEnv(VectorEnv): self.action_space = self.envs[0].action_space self.observation_space = self.envs[0].observation_space + @override(VectorEnv) def vector_reset(self): return [e.reset() for e in self.envs] + @override(VectorEnv) def reset_at(self, index): return self.envs[index].reset() + @override(VectorEnv) def vector_step(self, actions): obs_batch, rew_batch, done_batch, info_batch = [], [], [], [] for i in range(self.num_envs): @@ -88,5 +93,6 @@ class _VectorizedGymEnv(VectorEnv): info_batch.append(info) return obs_batch, rew_batch, done_batch, info_batch + @override(VectorEnv) def get_unwrapped(self): return self.envs diff --git a/python/ray/rllib/evaluation/policy_evaluator.py b/python/ray/rllib/evaluation/policy_evaluator.py index 360139839..b97f6a27b 100644 --- a/python/ray/rllib/evaluation/policy_evaluator.py +++ b/python/ray/rllib/evaluation/policy_evaluator.py @@ -21,6 +21,7 @@ from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph from ray.rllib.models import ModelCatalog from ray.rllib.models.preprocessors import NoPreprocessor from ray.rllib.utils import merge_dicts +from ray.rllib.utils.annotations import override from ray.rllib.utils.compression import pack from ray.rllib.utils.filter import get_filter from ray.rllib.utils.tf_run_builder import TFRunBuilder @@ -311,29 +312,7 @@ class PolicyEvaluator(EvaluatorInterface): logger.debug("Created evaluator with env {} ({}), policies {}".format( self.async_env, self.env, self.policy_map)) - def _build_policy_map(self, policy_dict, policy_config): - policy_map = {} - preprocessors = {} - for name, (cls, obs_space, act_space, - conf) in sorted(policy_dict.items()): - merged_conf = merge_dicts(policy_config, conf) - if self.preprocessing_enabled: - preprocessor = ModelCatalog.get_preprocessor_for_space( - obs_space, merged_conf.get("model")) - preprocessors[name] = preprocessor - obs_space = preprocessor.observation_space - else: - preprocessors[name] = NoPreprocessor(obs_space) - if isinstance(obs_space, gym.spaces.Dict) or \ - isinstance(obs_space, gym.spaces.Tuple): - raise ValueError( - "Found raw Tuple|Dict space as input to policy graph. " - "Please preprocess these observations with a " - "Tuple|DictFlatteningPreprocessor.") - with tf.variable_scope(name): - policy_map[name] = cls(obs_space, act_space, merged_conf) - return policy_map, preprocessors - + @override(EvaluatorInterface) def sample(self): """Evaluate the current policies and return a batch of experiences. @@ -382,6 +361,90 @@ class PolicyEvaluator(EvaluatorInterface): batch = self.sample() return batch, batch.count + @override(EvaluatorInterface) + def get_weights(self, policies=None): + if policies is None: + policies = self.policy_map.keys() + return { + pid: policy.get_weights() + for pid, policy in self.policy_map.items() if pid in policies + } + + @override(EvaluatorInterface) + def set_weights(self, weights): + for pid, w in weights.items(): + self.policy_map[pid].set_weights(w) + + @override(EvaluatorInterface) + def compute_gradients(self, samples): + if isinstance(samples, MultiAgentBatch): + grad_out, info_out = {}, {} + if self.tf_sess is not None: + builder = TFRunBuilder(self.tf_sess, "compute_gradients") + for pid, batch in samples.policy_batches.items(): + if pid not in self.policies_to_train: + continue + grad_out[pid], info_out[pid] = ( + self.policy_map[pid]._build_compute_gradients( + builder, batch)) + grad_out = {k: builder.get(v) for k, v in grad_out.items()} + info_out = {k: builder.get(v) for k, v in info_out.items()} + else: + for pid, batch in samples.policy_batches.items(): + if pid not in self.policies_to_train: + continue + grad_out[pid], info_out[pid] = ( + self.policy_map[pid].compute_gradients(batch)) + else: + grad_out, info_out = ( + self.policy_map[DEFAULT_POLICY_ID].compute_gradients(samples)) + info_out["batch_count"] = samples.count + return grad_out, info_out + + @override(EvaluatorInterface) + def apply_gradients(self, grads): + if isinstance(grads, dict): + if self.tf_sess is not None: + builder = TFRunBuilder(self.tf_sess, "apply_gradients") + outputs = { + pid: self.policy_map[pid]._build_apply_gradients( + builder, grad) + for pid, grad in grads.items() + } + return {k: builder.get(v) for k, v in outputs.items()} + else: + return { + pid: self.policy_map[pid].apply_gradients(g) + for pid, g in grads.items() + } + else: + return self.policy_map[DEFAULT_POLICY_ID].apply_gradients(grads) + + @override(EvaluatorInterface) + def compute_apply(self, samples): + if isinstance(samples, MultiAgentBatch): + info_out = {} + if self.tf_sess is not None: + builder = TFRunBuilder(self.tf_sess, "compute_apply") + for pid, batch in samples.policy_batches.items(): + if pid not in self.policies_to_train: + continue + info_out[pid], _ = ( + self.policy_map[pid]._build_compute_apply( + builder, batch)) + info_out = {k: builder.get(v) for k, v in info_out.items()} + else: + for pid, batch in samples.policy_batches.items(): + if pid not in self.policies_to_train: + continue + info_out[pid], _ = ( + self.policy_map[pid].compute_apply(batch)) + return info_out + else: + grad_fetch, apply_fetch = ( + self.policy_map[DEFAULT_POLICY_ID].compute_apply(samples)) + return grad_fetch + def for_policy(self, func, policy_id=DEFAULT_POLICY_ID): """Apply the given function to the specified policy graph.""" @@ -428,85 +491,6 @@ class PolicyEvaluator(EvaluatorInterface): f.clear_buffer() return return_filters - def get_weights(self, policies=None): - if policies is None: - policies = self.policy_map.keys() - return { - pid: policy.get_weights() - for pid, policy in self.policy_map.items() if pid in policies - } - - def set_weights(self, weights): - for pid, w in weights.items(): - self.policy_map[pid].set_weights(w) - - def compute_gradients(self, samples): - if isinstance(samples, MultiAgentBatch): - grad_out, info_out = {}, {} - if self.tf_sess is not None: - builder = TFRunBuilder(self.tf_sess, "compute_gradients") - for pid, batch in samples.policy_batches.items(): - if pid not in self.policies_to_train: - continue - grad_out[pid], info_out[pid] = ( - self.policy_map[pid].build_compute_gradients( - builder, batch)) - grad_out = {k: builder.get(v) for k, v in grad_out.items()} - info_out = {k: builder.get(v) for k, v in info_out.items()} - else: - for pid, batch in samples.policy_batches.items(): - if pid not in self.policies_to_train: - continue - grad_out[pid], info_out[pid] = ( - self.policy_map[pid].compute_gradients(batch)) - else: - grad_out, info_out = ( - self.policy_map[DEFAULT_POLICY_ID].compute_gradients(samples)) - info_out["batch_count"] = samples.count - return grad_out, info_out - - def apply_gradients(self, grads): - if isinstance(grads, dict): - if self.tf_sess is not None: - builder = TFRunBuilder(self.tf_sess, "apply_gradients") - outputs = { - pid: self.policy_map[pid].build_apply_gradients( - builder, grad) - for pid, grad in grads.items() - } - return {k: builder.get(v) for k, v in outputs.items()} - else: - return { - pid: self.policy_map[pid].apply_gradients(g) - for pid, g in grads.items() - } - else: - return self.policy_map[DEFAULT_POLICY_ID].apply_gradients(grads) - - def compute_apply(self, samples): - if isinstance(samples, MultiAgentBatch): - info_out = {} - if self.tf_sess is not None: - builder = TFRunBuilder(self.tf_sess, "compute_apply") - for pid, batch in samples.policy_batches.items(): - if pid not in self.policies_to_train: - continue - info_out[pid], _ = ( - self.policy_map[pid].build_compute_apply( - builder, batch)) - info_out = {k: builder.get(v) for k, v in info_out.items()} - else: - for pid, batch in samples.policy_batches.items(): - if pid not in self.policies_to_train: - continue - info_out[pid], _ = ( - self.policy_map[pid].compute_apply(batch)) - return info_out - else: - grad_fetch, apply_fetch = ( - self.policy_map[DEFAULT_POLICY_ID].compute_apply(samples)) - return grad_fetch - def save(self): filters = self.get_filters(flush_after=True) state = { @@ -524,6 +508,29 @@ class PolicyEvaluator(EvaluatorInterface): def set_global_vars(self, global_vars): self.foreach_policy(lambda p, _: p.on_global_var_update(global_vars)) + def _build_policy_map(self, policy_dict, policy_config): + policy_map = {} + preprocessors = {} + for name, (cls, obs_space, act_space, + conf) in sorted(policy_dict.items()): + merged_conf = merge_dicts(policy_config, conf) + if self.preprocessing_enabled: + preprocessor = ModelCatalog.get_preprocessor_for_space( + obs_space, merged_conf.get("model")) + preprocessors[name] = preprocessor + obs_space = preprocessor.observation_space + else: + preprocessors[name] = NoPreprocessor(obs_space) + if isinstance(obs_space, gym.spaces.Dict) or \ + isinstance(obs_space, gym.spaces.Tuple): + raise ValueError( + "Found raw Tuple|Dict space as input to policy graph. " + "Please preprocess these observations with a " + "Tuple|DictFlatteningPreprocessor.") + with tf.variable_scope(name): + policy_map[name] = cls(obs_space, act_space, merged_conf) + return policy_map, preprocessors + def _validate_and_canonicalize(policy_graph, env): if isinstance(policy_graph, dict): diff --git a/python/ray/rllib/evaluation/sampler.py b/python/ray/rllib/evaluation/sampler.py index cc9e32295..ac7c6ed8a 100644 --- a/python/ray/rllib/evaluation/sampler.py +++ b/python/ray/rllib/evaluation/sampler.py @@ -467,7 +467,7 @@ def _do_policy_eval(tf_sess, to_eval, policies, active_episodes, clip_actions): policy = _get_or_raise(policies, policy_id) if builder and (policy.compute_actions.__code__ is TFPolicyGraph.compute_actions.__code__): - pending_fetches[policy_id] = policy.build_compute_actions( + pending_fetches[policy_id] = policy._build_compute_actions( builder, [t.obs for t in eval_data], rnn_in_cols, prev_action_batch=[t.prev_action for t in eval_data], diff --git a/python/ray/rllib/evaluation/tf_policy_graph.py b/python/ray/rllib/evaluation/tf_policy_graph.py index 95e7a5d66..e5a1d7b19 100644 --- a/python/ray/rllib/evaluation/tf_policy_graph.py +++ b/python/ray/rllib/evaluation/tf_policy_graph.py @@ -9,8 +9,9 @@ import numpy as np import ray from ray.rllib.evaluation.policy_graph import PolicyGraph from ray.rllib.models.lstm import chop_into_sequences -from ray.rllib.utils.tf_run_builder import TFRunBuilder +from ray.rllib.utils.annotations import override from ray.rllib.utils.schedules import ConstantSchedule, PiecewiseSchedule +from ray.rllib.utils.tf_run_builder import TFRunBuilder logger = logging.getLogger(__name__) @@ -146,13 +147,90 @@ class TFPolicyGraph(PolicyGraph): logger.debug("Created {} with loss inputs: {}".format( self, self._loss_input_dict)) - def build_compute_actions(self, - builder, - obs_batch, - state_batches=None, - prev_action_batch=None, - prev_reward_batch=None, - episodes=None): + @override(PolicyGraph) + def compute_actions(self, + obs_batch, + state_batches=None, + prev_action_batch=None, + prev_reward_batch=None, + episodes=None): + builder = TFRunBuilder(self._sess, "compute_actions") + fetches = self._build_compute_actions(builder, obs_batch, + state_batches, prev_action_batch, + prev_reward_batch) + return builder.get(fetches) + + @override(PolicyGraph) + def compute_gradients(self, postprocessed_batch): + builder = TFRunBuilder(self._sess, "compute_gradients") + fetches = self._build_compute_gradients(builder, postprocessed_batch) + return builder.get(fetches) + + @override(PolicyGraph) + def apply_gradients(self, gradients): + builder = TFRunBuilder(self._sess, "apply_gradients") + fetches = self._build_apply_gradients(builder, gradients) + return builder.get(fetches) + + @override(PolicyGraph) + def compute_apply(self, postprocessed_batch): + builder = TFRunBuilder(self._sess, "compute_apply") + fetches = self._build_compute_apply(builder, postprocessed_batch) + return builder.get(fetches) + + @override(PolicyGraph) + def get_weights(self): + return self._variables.get_flat() + + @override(PolicyGraph) + def set_weights(self, weights): + return self._variables.set_flat(weights) + + def copy(self, existing_inputs): + """Creates a copy of self using existing input placeholders. + + Optional, only required to work with the multi-GPU optimizer.""" + raise NotImplementedError + + def extra_compute_action_feed_dict(self): + """Extra dict to pass to the compute actions session run.""" + return {} + + def extra_compute_action_fetches(self): + """Extra values to fetch and return from compute_actions().""" + return {} # e.g, value function + + def extra_compute_grad_feed_dict(self): + """Extra dict to pass to the compute gradients session run.""" + return {} # e.g, kl_coeff + + def extra_compute_grad_fetches(self): + """Extra values to fetch and return from compute_gradients().""" + return {} # e.g, td error + + def extra_apply_grad_feed_dict(self): + """Extra dict to pass to the apply gradients session run.""" + return {} + + def extra_apply_grad_fetches(self): + """Extra values to fetch and return from apply_gradients().""" + return {} # e.g., batch norm updates + + def optimizer(self): + """TF optimizer to use for policy optimization.""" + return tf.train.AdamOptimizer() + + def gradients(self, optimizer): + """Override for custom gradient computation.""" + return optimizer.compute_gradients(self._loss) + + def _build_compute_actions(self, + builder, + obs_batch, + state_batches=None, + prev_action_batch=None, + prev_reward_batch=None, + episodes=None): state_batches = state_batches or [] assert len(self._state_inputs) == len(state_batches), \ (self._state_inputs, state_batches) @@ -170,17 +248,43 @@ class TFPolicyGraph(PolicyGraph): [self.extra_compute_action_fetches()]) return fetches[0], fetches[1:-1], fetches[-1] - def compute_actions(self, - obs_batch, - state_batches=None, - prev_action_batch=None, - prev_reward_batch=None, - episodes=None): - builder = TFRunBuilder(self._sess, "compute_actions") - fetches = self.build_compute_actions(builder, obs_batch, state_batches, - prev_action_batch, - prev_reward_batch) - return builder.get(fetches) + def _build_compute_gradients(self, builder, postprocessed_batch): + builder.add_feed_dict(self.extra_compute_grad_feed_dict()) + builder.add_feed_dict({self._is_training: True}) + builder.add_feed_dict(self._get_loss_inputs_dict(postprocessed_batch)) + fetches = builder.add_fetches( + [self._grads, self.extra_compute_grad_fetches()]) + return fetches[0], fetches[1] + + def _build_apply_gradients(self, builder, gradients): + assert len(gradients) == len(self._grads), (gradients, self._grads) + builder.add_feed_dict(self.extra_apply_grad_feed_dict()) + builder.add_feed_dict({self._is_training: True}) + builder.add_feed_dict(dict(zip(self._grads, gradients))) + fetches = builder.add_fetches( + [self._apply_op, self.extra_apply_grad_fetches()]) + return fetches[1] + + def _build_compute_apply(self, builder, postprocessed_batch): + builder.add_feed_dict(self.extra_compute_grad_feed_dict()) + builder.add_feed_dict(self.extra_apply_grad_feed_dict()) + builder.add_feed_dict(self._get_loss_inputs_dict(postprocessed_batch)) + builder.add_feed_dict({self._is_training: True}) + fetches = builder.add_fetches([ + self._apply_op, + self.extra_compute_grad_fetches(), + self.extra_apply_grad_fetches() + ]) + return fetches[1], fetches[2] + + def _get_is_training_placeholder(self): + """Get the placeholder for _is_training, i.e., for batch norm layers. + + This can be called safely before __init__ has run. + """ + if not hasattr(self, "_is_training"): + self._is_training = tf.placeholder_with_default(False, ()) + return self._is_training def _get_loss_inputs_dict(self, batch): feed_dict = {} @@ -222,92 +326,6 @@ class TFPolicyGraph(PolicyGraph): feed_dict[self._seq_lens] = seq_lens return feed_dict - def build_compute_gradients(self, builder, postprocessed_batch): - builder.add_feed_dict(self.extra_compute_grad_feed_dict()) - builder.add_feed_dict({self._is_training: True}) - builder.add_feed_dict(self._get_loss_inputs_dict(postprocessed_batch)) - fetches = builder.add_fetches( - [self._grads, self.extra_compute_grad_fetches()]) - return fetches[0], fetches[1] - - def compute_gradients(self, postprocessed_batch): - builder = TFRunBuilder(self._sess, "compute_gradients") - fetches = self.build_compute_gradients(builder, postprocessed_batch) - return builder.get(fetches) - - def build_apply_gradients(self, builder, gradients): - assert len(gradients) == len(self._grads), (gradients, self._grads) - builder.add_feed_dict(self.extra_apply_grad_feed_dict()) - builder.add_feed_dict({self._is_training: True}) - builder.add_feed_dict(dict(zip(self._grads, gradients))) - fetches = builder.add_fetches( - [self._apply_op, self.extra_apply_grad_fetches()]) - return fetches[1] - - def apply_gradients(self, gradients): - builder = TFRunBuilder(self._sess, "apply_gradients") - fetches = self.build_apply_gradients(builder, gradients) - return builder.get(fetches) - - def build_compute_apply(self, builder, postprocessed_batch): - builder.add_feed_dict(self.extra_compute_grad_feed_dict()) - builder.add_feed_dict(self.extra_apply_grad_feed_dict()) - builder.add_feed_dict(self._get_loss_inputs_dict(postprocessed_batch)) - builder.add_feed_dict({self._is_training: True}) - fetches = builder.add_fetches([ - self._apply_op, - self.extra_compute_grad_fetches(), - self.extra_apply_grad_fetches() - ]) - return fetches[1], fetches[2] - - def compute_apply(self, postprocessed_batch): - builder = TFRunBuilder(self._sess, "compute_apply") - fetches = self.build_compute_apply(builder, postprocessed_batch) - return builder.get(fetches) - - def get_weights(self): - return self._variables.get_flat() - - def set_weights(self, weights): - return self._variables.set_flat(weights) - - def extra_compute_action_feed_dict(self): - return {} - - def extra_compute_action_fetches(self): - return {} # e.g, value function - - def extra_compute_grad_feed_dict(self): - return {} # e.g, kl_coeff - - def extra_compute_grad_fetches(self): - return {} # e.g, td error - - def extra_apply_grad_feed_dict(self): - return {} - - def extra_apply_grad_fetches(self): - return {} # e.g., batch norm updates - - def optimizer(self): - return tf.train.AdamOptimizer() - - def gradients(self, optimizer): - return optimizer.compute_gradients(self._loss) - - def loss_inputs(self): - return self._loss_inputs - - def _get_is_training_placeholder(self): - """Get the placeholder for _is_training, i.e., for batch norm layers. - - This can be called safely before __init__ has run. - """ - if not hasattr(self, "_is_training"): - self._is_training = tf.placeholder_with_default(False, ()) - return self._is_training - class LearningRateSchedule(object): """Mixin for TFPolicyGraph that adds a learning rate schedule.""" @@ -320,11 +338,13 @@ class LearningRateSchedule(object): self.lr_schedule = PiecewiseSchedule( lr_schedule, outside_value=lr_schedule[-1][-1]) + @override(PolicyGraph) def on_global_var_update(self, global_vars): super(LearningRateSchedule, self).on_global_var_update(global_vars) self.cur_lr.load( self.lr_schedule.value(global_vars["timestep"]), session=self._sess) + @override(TFPolicyGraph) def optimizer(self): return tf.train.AdamOptimizer(self.cur_lr) diff --git a/python/ray/rllib/evaluation/torch_policy_graph.py b/python/ray/rllib/evaluation/torch_policy_graph.py index a762927ba..c8e86e845 100644 --- a/python/ray/rllib/evaluation/torch_policy_graph.py +++ b/python/ray/rllib/evaluation/torch_policy_graph.py @@ -13,6 +13,7 @@ except ImportError: pass # soft dep from ray.rllib.evaluation.policy_graph import PolicyGraph +from ray.rllib.utils.annotations import override class TorchPolicyGraph(PolicyGraph): @@ -56,17 +57,7 @@ class TorchPolicyGraph(PolicyGraph): self._loss_inputs = loss_inputs self._optimizer = self.optimizer() - def extra_action_out(self, model_out): - """Returns dict of extra info to include in experience batch. - - Arguments: - model_out (list): Outputs of the policy model module.""" - return {} - - def optimizer(self): - """Custom PyTorch optimizer to use.""" - return torch.optim.Adam(self._model.parameters()) - + @override(PolicyGraph) def compute_actions(self, obs_batch, state_batches=None, @@ -83,6 +74,7 @@ class TorchPolicyGraph(PolicyGraph): actions = F.softmax(logits, dim=1).multinomial(1).squeeze(0) return var_to_np(actions), [], self.extra_action_out(model_out) + @override(PolicyGraph) def compute_gradients(self, postprocessed_batch): with self.lock: loss_in = [] @@ -96,6 +88,7 @@ class TorchPolicyGraph(PolicyGraph): grads = [var_to_np(p.grad.data) for p in self._model.parameters()] return grads, {} + @override(PolicyGraph) def apply_gradients(self, gradients): with self.lock: for g, p in zip(gradients, self._model.parameters()): @@ -103,10 +96,23 @@ class TorchPolicyGraph(PolicyGraph): self._optimizer.step() return {} + @override(PolicyGraph) def get_weights(self): with self.lock: return self._model.state_dict() + @override(PolicyGraph) def set_weights(self, weights): with self.lock: self._model.load_state_dict(weights) + + def extra_action_out(self, model_out): + """Returns dict of extra info to include in experience batch. + + Arguments: + model_out (list): Outputs of the policy model module.""" + return {} + + def optimizer(self): + """Custom PyTorch optimizer to use.""" + return torch.optim.Adam(self._model.parameters()) diff --git a/python/ray/rllib/models/action_dist.py b/python/ray/rllib/models/action_dist.py index 76d45e244..f2a69efaf 100644 --- a/python/ray/rllib/models/action_dist.py +++ b/python/ray/rllib/models/action_dist.py @@ -4,10 +4,11 @@ from __future__ import print_function from collections import namedtuple import distutils.version - import tensorflow as tf import numpy as np +from ray.rllib.utils.annotations import override + use_tf150_api = (distutils.version.LooseVersion(tf.VERSION) >= distutils.version.LooseVersion("1.5.0")) @@ -42,10 +43,12 @@ class ActionDistribution(object): class Categorical(ActionDistribution): """Categorical distribution for discrete action spaces.""" + @override(ActionDistribution) def logp(self, x): return -tf.nn.sparse_softmax_cross_entropy_with_logits( logits=self.inputs, labels=x) + @override(ActionDistribution) def entropy(self): if use_tf150_api: a0 = self.inputs - tf.reduce_max( @@ -61,6 +64,7 @@ class Categorical(ActionDistribution): p0 = ea0 / z0 return tf.reduce_sum(p0 * (tf.log(z0) - a0), reduction_indices=[1]) + @override(ActionDistribution) def kl(self, other): if use_tf150_api: a0 = self.inputs - tf.reduce_max( @@ -84,6 +88,7 @@ class Categorical(ActionDistribution): return tf.reduce_sum( p0 * (a0 - tf.log(z0) - a1 + tf.log(z1)), reduction_indices=[1]) + @override(ActionDistribution) def sample(self): return tf.squeeze(tf.multinomial(self.inputs, 1), axis=1) @@ -102,12 +107,14 @@ class DiagGaussian(ActionDistribution): self.log_std = log_std self.std = tf.exp(log_std) + @override(ActionDistribution) def logp(self, x): return (-0.5 * tf.reduce_sum( tf.square((x - self.mean) / self.std), reduction_indices=[1]) - 0.5 * np.log(2.0 * np.pi) * tf.to_float(tf.shape(x)[1]) - tf.reduce_sum(self.log_std, reduction_indices=[1])) + @override(ActionDistribution) def kl(self, other): assert isinstance(other, DiagGaussian) return tf.reduce_sum( @@ -116,11 +123,13 @@ class DiagGaussian(ActionDistribution): (2.0 * tf.square(other.std)) - 0.5, reduction_indices=[1]) + @override(ActionDistribution) def entropy(self): return tf.reduce_sum( .5 * self.log_std + .5 * np.log(2.0 * np.pi * np.e), reduction_indices=[1]) + @override(ActionDistribution) def sample(self): return self.mean + self.std * tf.random_normal(tf.shape(self.mean)) @@ -131,6 +140,7 @@ class Deterministic(ActionDistribution): This is similar to DiagGaussian with standard deviation zero. """ + @override(ActionDistribution) def sample(self): return self.inputs @@ -150,8 +160,8 @@ class MultiActionDistribution(ActionDistribution): child_list.append(distribution(split_inputs[i])) self.child_distributions = child_list + @override(ActionDistribution) def logp(self, x): - """The log-likelihood of the action distribution.""" split_indices = [] for dist in self.child_distributions: if isinstance(dist, Categorical): @@ -170,8 +180,8 @@ class MultiActionDistribution(ActionDistribution): ]) return np.sum(log_list) + @override(ActionDistribution) def kl(self, other): - """The KL-divergence between two action distributions.""" kl_list = np.asarray([ distribution.kl(other_distribution) for distribution, other_distribution in zip( @@ -179,15 +189,14 @@ class MultiActionDistribution(ActionDistribution): ]) return np.sum(kl_list) + @override(ActionDistribution) def entropy(self): - """The entropy of the action distribution.""" entropy_list = np.array( [s.entropy() for s in self.child_distributions]) return np.sum(entropy_list) + @override(ActionDistribution) def sample(self): - """Draw a sample from the action distribution.""" - return TupleActions([s.sample() for s in self.child_distributions]) diff --git a/python/ray/rllib/models/fcnet.py b/python/ray/rllib/models/fcnet.py index 5a759fd59..19745b9e7 100644 --- a/python/ray/rllib/models/fcnet.py +++ b/python/ray/rllib/models/fcnet.py @@ -7,11 +7,13 @@ import tensorflow.contrib.slim as slim from ray.rllib.models.model import Model from ray.rllib.models.misc import normc_initializer, get_activation_fn +from ray.rllib.utils.annotations import override class FullyConnectedNetwork(Model): """Generic fully connected network.""" + @override(Model) def _build_layers(self, inputs, num_outputs, options): """Process the flattened inputs. diff --git a/python/ray/rllib/models/lstm.py b/python/ray/rllib/models/lstm.py index fdb7af602..323c7f375 100644 --- a/python/ray/rllib/models/lstm.py +++ b/python/ray/rllib/models/lstm.py @@ -23,6 +23,72 @@ import tensorflow.contrib.rnn as rnn from ray.rllib.models.misc import linear, normc_initializer from ray.rllib.models.model import Model +from ray.rllib.utils.annotations import override + + +class LSTM(Model): + """Adds a LSTM cell on top of some other model output. + + Uses a linear layer at the end for output. + + Important: we assume inputs is a padded batch of sequences denoted by + self.seq_lens. See add_time_dimension() for more information. + """ + + @override(Model) + def _build_layers_v2(self, input_dict, num_outputs, options): + cell_size = options.get("lstm_cell_size") + if options.get("lstm_use_prev_action_reward"): + action_dim = int( + np.product( + input_dict["prev_actions"].get_shape().as_list()[1:])) + features = tf.concat( + [ + input_dict["obs"], + tf.reshape( + tf.cast(input_dict["prev_actions"], tf.float32), + [-1, action_dim]), + tf.reshape(input_dict["prev_rewards"], [-1, 1]), + ], + axis=1) + else: + features = input_dict["obs"] + last_layer = add_time_dimension(features, self.seq_lens) + + # Setup the LSTM cell + lstm = rnn.BasicLSTMCell(cell_size, state_is_tuple=True) + self.state_init = [ + np.zeros(lstm.state_size.c, np.float32), + np.zeros(lstm.state_size.h, np.float32) + ] + + # Setup LSTM inputs + if self.state_in: + c_in, h_in = self.state_in + else: + c_in = tf.placeholder( + tf.float32, [None, lstm.state_size.c], name="c") + h_in = tf.placeholder( + tf.float32, [None, lstm.state_size.h], name="h") + self.state_in = [c_in, h_in] + + # Setup LSTM outputs + state_in = rnn.LSTMStateTuple(c_in, h_in) + lstm_out, lstm_state = tf.nn.dynamic_rnn( + lstm, + last_layer, + initial_state=state_in, + sequence_length=self.seq_lens, + time_major=False, + dtype=tf.float32) + + self.state_out = list(lstm_state) + + # Compute outputs + last_layer = tf.reshape(lstm_out, [-1, cell_size]) + logits = linear(last_layer, num_outputs, "action", + normc_initializer(0.01)) + return logits, last_layer def add_time_dimension(padded_inputs, seq_lens): @@ -138,67 +204,3 @@ def chop_into_sequences(episode_ids, initial_states.append(np.array(s_init)) return feature_sequences, initial_states, np.array(seq_lens) - - -class LSTM(Model): - """Adds a LSTM cell on top of some other model output. - - Uses a linear layer at the end for output. - - Important: we assume inputs is a padded batch of sequences denoted by - self.seq_lens. See add_time_dimension() for more information. - """ - - def _build_layers_v2(self, input_dict, num_outputs, options): - cell_size = options.get("lstm_cell_size") - if options.get("lstm_use_prev_action_reward"): - action_dim = int( - np.product( - input_dict["prev_actions"].get_shape().as_list()[1:])) - features = tf.concat( - [ - input_dict["obs"], - tf.reshape( - tf.cast(input_dict["prev_actions"], tf.float32), - [-1, action_dim]), - tf.reshape(input_dict["prev_rewards"], [-1, 1]), - ], - axis=1) - else: - features = input_dict["obs"] - last_layer = add_time_dimension(features, self.seq_lens) - - # Setup the LSTM cell - lstm = rnn.BasicLSTMCell(cell_size, state_is_tuple=True) - self.state_init = [ - np.zeros(lstm.state_size.c, np.float32), - np.zeros(lstm.state_size.h, np.float32) - ] - - # Setup LSTM inputs - if self.state_in: - c_in, h_in = self.state_in - else: - c_in = tf.placeholder( - tf.float32, [None, lstm.state_size.c], name="c") - h_in = tf.placeholder( - tf.float32, [None, lstm.state_size.h], name="h") - self.state_in = [c_in, h_in] - - # Setup LSTM outputs - state_in = rnn.LSTMStateTuple(c_in, h_in) - lstm_out, lstm_state = tf.nn.dynamic_rnn( - lstm, - last_layer, - initial_state=state_in, - sequence_length=self.seq_lens, - time_major=False, - dtype=tf.float32) - - self.state_out = list(lstm_state) - - # Compute outputs - last_layer = tf.reshape(lstm_out, [-1, cell_size]) - logits = linear(last_layer, num_outputs, "action", - normc_initializer(0.01)) - return logits, last_layer diff --git a/python/ray/rllib/models/model.py b/python/ray/rllib/models/model.py index 561b636dc..818966bb1 100644 --- a/python/ray/rllib/models/model.py +++ b/python/ray/rllib/models/model.py @@ -82,19 +82,6 @@ class Model(object): self.outputs = tf.concat( [self.outputs, 0.0 * self.outputs + log_std], 1) - def _validate_output_shape(self): - """Checks that the model has the correct number of outputs.""" - try: - out = tf.convert_to_tensor(self.outputs) - shape = out.shape.as_list() - except Exception: - raise ValueError("Output is not a tensor: {}".format(self.outputs)) - else: - if len(shape) != 2 or shape[1] != self._num_outputs: - raise ValueError( - "Expected output shape of [None, {}], got {}".format( - self._num_outputs, shape)) - def _build_layers(self, inputs, num_outputs, options): """Builds and returns the output and last layer of the network. @@ -159,6 +146,19 @@ class Model(object): """ return tf.constant(0.0) + def _validate_output_shape(self): + """Checks that the model has the correct number of outputs.""" + try: + out = tf.convert_to_tensor(self.outputs) + shape = out.shape.as_list() + except Exception: + raise ValueError("Output is not a tensor: {}".format(self.outputs)) + else: + if len(shape) != 2 or shape[1] != self._num_outputs: + raise ValueError( + "Expected output shape of [None, {}], got {}".format( + self._num_outputs, shape)) + def _restore_original_dimensions(input_dict, obs_space): if hasattr(obs_space, "original_space"): diff --git a/python/ray/rllib/models/preprocessors.py b/python/ray/rllib/models/preprocessors.py index 66835e864..0238ef2d8 100644 --- a/python/ray/rllib/models/preprocessors.py +++ b/python/ray/rllib/models/preprocessors.py @@ -8,6 +8,8 @@ import logging import numpy as np import gym +from ray.rllib.utils.annotations import override + ATARI_OBS_SHAPE = (210, 160, 3) ATARI_RAM_OBS_SHAPE = (128, ) @@ -57,6 +59,7 @@ class GenericPixelPreprocessor(Preprocessor): instead for deepmind-style Atari preprocessing. """ + @override(Preprocessor) def _init_shape(self, obs_space, options): self._grayscale = options.get("grayscale") self._zero_mean = options.get("zero_mean") @@ -72,6 +75,7 @@ class GenericPixelPreprocessor(Preprocessor): shape = shape[-1:] + shape[:-1] return shape + @override(Preprocessor) def transform(self, observation): """Downsamples images from (210, 160, 3) by the configured factor.""" scaled = observation[25:-25, :, :] @@ -96,17 +100,21 @@ class GenericPixelPreprocessor(Preprocessor): class AtariRamPreprocessor(Preprocessor): + @override(Preprocessor) def _init_shape(self, obs_space, options): return (128, ) + @override(Preprocessor) def transform(self, observation): return (observation - 128) / 128 class OneHotPreprocessor(Preprocessor): + @override(Preprocessor) def _init_shape(self, obs_space, options): return (self._obs_space.n, ) + @override(Preprocessor) def transform(self, observation): arr = np.zeros(self._obs_space.n) if not self._obs_space.contains(observation): @@ -117,9 +125,11 @@ class OneHotPreprocessor(Preprocessor): class NoPreprocessor(Preprocessor): + @override(Preprocessor) def _init_shape(self, obs_space, options): return self._obs_space.shape + @override(Preprocessor) def transform(self, observation): return observation @@ -130,6 +140,7 @@ class TupleFlatteningPreprocessor(Preprocessor): RLlib models will unpack the flattened output before _build_layers_v2(). """ + @override(Preprocessor) def _init_shape(self, obs_space, options): assert isinstance(self._obs_space, gym.spaces.Tuple) size = 0 @@ -142,6 +153,7 @@ class TupleFlatteningPreprocessor(Preprocessor): size += preprocessor.size return (size, ) + @override(Preprocessor) def transform(self, observation): assert len(observation) == len(self.preprocessors), observation return np.concatenate([ @@ -156,6 +168,7 @@ class DictFlatteningPreprocessor(Preprocessor): RLlib models will unpack the flattened output before _build_layers_v2(). """ + @override(Preprocessor) def _init_shape(self, obs_space, options): assert isinstance(self._obs_space, gym.spaces.Dict) size = 0 @@ -167,6 +180,7 @@ class DictFlatteningPreprocessor(Preprocessor): size += preprocessor.size return (size, ) + @override(Preprocessor) def transform(self, observation): if not isinstance(observation, OrderedDict): observation = OrderedDict(sorted(list(observation.items()))) diff --git a/python/ray/rllib/models/visionnet.py b/python/ray/rllib/models/visionnet.py index 1d856e42c..0638c4fc8 100644 --- a/python/ray/rllib/models/visionnet.py +++ b/python/ray/rllib/models/visionnet.py @@ -7,16 +7,18 @@ import tensorflow.contrib.slim as slim from ray.rllib.models.model import Model from ray.rllib.models.misc import get_activation_fn, flatten +from ray.rllib.utils.annotations import override class VisionNetwork(Model): """Generic vision network.""" + @override(Model) def _build_layers_v2(self, input_dict, num_outputs, options): inputs = input_dict["obs"] filters = options.get("conv_filters") if not filters: - filters = get_filter_config(inputs) + filters = _get_filter_config(inputs) activation = get_activation_fn(options.get("conv_activation")) @@ -47,7 +49,7 @@ class VisionNetwork(Model): return flatten(fc2), flatten(fc1) -def get_filter_config(inputs): +def _get_filter_config(inputs): filters_84x84 = [ [16, [8, 8], 4], [32, [4, 4], 2], diff --git a/python/ray/rllib/optimizers/async_gradients_optimizer.py b/python/ray/rllib/optimizers/async_gradients_optimizer.py index 499d2a91f..b1e5ebe84 100644 --- a/python/ray/rllib/optimizers/async_gradients_optimizer.py +++ b/python/ray/rllib/optimizers/async_gradients_optimizer.py @@ -4,6 +4,7 @@ from __future__ import print_function import ray from ray.rllib.optimizers.policy_optimizer import PolicyOptimizer +from ray.rllib.utils.annotations import override from ray.rllib.utils.timer import TimerStat @@ -15,6 +16,7 @@ class AsyncGradientsOptimizer(PolicyOptimizer): gradient computations on the remote workers. """ + @override(PolicyOptimizer) def _init(self, grads_per_step=100): self.apply_timer = TimerStat() self.wait_timer = TimerStat() @@ -25,6 +27,7 @@ class AsyncGradientsOptimizer(PolicyOptimizer): raise ValueError( "Async optimizer requires at least 1 remote evaluator") + @override(PolicyOptimizer) def step(self): weights = ray.put(self.local_evaluator.get_weights()) pending_gradients = {} @@ -64,6 +67,7 @@ class AsyncGradientsOptimizer(PolicyOptimizer): pending_gradients[future] = e num_gradients += 1 + @override(PolicyOptimizer) def stats(self): return dict( PolicyOptimizer.stats(self), **{ diff --git a/python/ray/rllib/optimizers/async_replay_optimizer.py b/python/ray/rllib/optimizers/async_replay_optimizer.py index 932800001..582bb6539 100644 --- a/python/ray/rllib/optimizers/async_replay_optimizer.py +++ b/python/ray/rllib/optimizers/async_replay_optimizer.py @@ -20,6 +20,7 @@ from ray.rllib.evaluation.sample_batch import SampleBatch, DEFAULT_POLICY_ID, \ MultiAgentBatch from ray.rllib.optimizers.policy_optimizer import PolicyOptimizer from ray.rllib.optimizers.replay_buffer import PrioritizedReplayBuffer +from ray.rllib.utils.annotations import override from ray.rllib.utils.actors import TaskPool, create_colocated from ray.rllib.utils.timer import TimerStat from ray.rllib.utils.window_stat import WindowStat @@ -29,6 +30,189 @@ REPLAY_QUEUE_DEPTH = 4 LEARNER_QUEUE_MAX_SIZE = 16 +class AsyncReplayOptimizer(PolicyOptimizer): + """Main event loop of the Ape-X optimizer (async sampling with replay). + + This class coordinates the data transfers between the learner thread, + remote evaluators (Ape-X actors), and replay buffer actors. + + This optimizer requires that policy evaluators return an additional + "td_error" array in the info return of compute_gradients(). This error + term will be used for sample prioritization.""" + + @override(PolicyOptimizer) + def _init(self, + learning_starts=1000, + buffer_size=10000, + prioritized_replay=True, + prioritized_replay_alpha=0.6, + prioritized_replay_beta=0.4, + prioritized_replay_eps=1e-6, + train_batch_size=512, + sample_batch_size=50, + num_replay_buffer_shards=1, + max_weight_sync_delay=400, + debug=False): + + self.debug = debug + self.replay_starts = learning_starts + self.prioritized_replay_beta = prioritized_replay_beta + self.prioritized_replay_eps = prioritized_replay_eps + self.max_weight_sync_delay = max_weight_sync_delay + + self.learner = LearnerThread(self.local_evaluator) + self.learner.start() + + self.replay_actors = create_colocated(ReplayActor, [ + num_replay_buffer_shards, + learning_starts, + buffer_size, + train_batch_size, + prioritized_replay_alpha, + prioritized_replay_beta, + prioritized_replay_eps, + ], num_replay_buffer_shards) + + # Stats + self.timers = { + k: TimerStat() + for k in [ + "put_weights", "get_samples", "sample_processing", + "replay_processing", "update_priorities", "train", "sample" + ] + } + self.num_weight_syncs = 0 + self.num_samples_dropped = 0 + self.learning_started = False + + # Number of worker steps since the last weight update + self.steps_since_update = {} + + # Otherwise kick of replay tasks for local gradient updates + self.replay_tasks = TaskPool() + for ra in self.replay_actors: + for _ in range(REPLAY_QUEUE_DEPTH): + self.replay_tasks.add(ra, ra.replay.remote()) + + # Kick off async background sampling + self.sample_tasks = TaskPool() + if self.remote_evaluators: + self._set_evaluators(self.remote_evaluators) + + @override(PolicyOptimizer) + def step(self): + assert len(self.remote_evaluators) > 0 + start = time.time() + sample_timesteps, train_timesteps = self._step() + time_delta = time.time() - start + self.timers["sample"].push(time_delta) + self.timers["sample"].push_units_processed(sample_timesteps) + if train_timesteps > 0: + self.learning_started = True + if self.learning_started: + self.timers["train"].push(time_delta) + self.timers["train"].push_units_processed(train_timesteps) + self.num_steps_sampled += sample_timesteps + self.num_steps_trained += train_timesteps + + @override(PolicyOptimizer) + def stop(self): + for r in self.replay_actors: + r.__ray_terminate__.remote() + self.learner.stopped = True + + @override(PolicyOptimizer) + def stats(self): + replay_stats = ray.get(self.replay_actors[0].stats.remote(self.debug)) + timing = { + "{}_time_ms".format(k): round(1000 * self.timers[k].mean, 3) + for k in self.timers + } + timing["learner_grad_time_ms"] = round( + 1000 * self.learner.grad_timer.mean, 3) + timing["learner_dequeue_time_ms"] = round( + 1000 * self.learner.queue_timer.mean, 3) + stats = { + "sample_throughput": round(self.timers["sample"].mean_throughput, + 3), + "train_throughput": round(self.timers["train"].mean_throughput, 3), + "num_weight_syncs": self.num_weight_syncs, + "num_samples_dropped": self.num_samples_dropped, + "learner_queue": self.learner.learner_queue_size.stats(), + "replay_shard_0": replay_stats, + } + debug_stats = { + "timing_breakdown": timing, + "pending_sample_tasks": self.sample_tasks.count, + "pending_replay_tasks": self.replay_tasks.count, + } + if self.debug: + stats.update(debug_stats) + if self.learner.stats: + stats["learner"] = self.learner.stats + return dict(PolicyOptimizer.stats(self), **stats) + + # For https://github.com/ray-project/ray/issues/2541 only + def _set_evaluators(self, remote_evaluators): + self.remote_evaluators = remote_evaluators + weights = self.local_evaluator.get_weights() + for ev in self.remote_evaluators: + ev.set_weights.remote(weights) + self.steps_since_update[ev] = 0 + for _ in range(SAMPLE_QUEUE_DEPTH): + self.sample_tasks.add(ev, ev.sample_with_count.remote()) + + def _step(self): + sample_timesteps, train_timesteps = 0, 0 + weights = None + + with self.timers["sample_processing"]: + completed = list(self.sample_tasks.completed()) + counts = ray.get([c[1][1] for c in completed]) + for i, (ev, (sample_batch, count)) in enumerate(completed): + sample_timesteps += counts[i] + + # Send the data to the replay buffer + random.choice( + self.replay_actors).add_batch.remote(sample_batch) + + # Update weights if needed + self.steps_since_update[ev] += counts[i] + if self.steps_since_update[ev] >= self.max_weight_sync_delay: + # Note that it's important to pull new weights once + # updated to avoid excessive correlation between actors + if weights is None or self.learner.weights_updated: + self.learner.weights_updated = False + with self.timers["put_weights"]: + weights = ray.put( + self.local_evaluator.get_weights()) + ev.set_weights.remote(weights) + self.num_weight_syncs += 1 + self.steps_since_update[ev] = 0 + + # Kick off another sample request + self.sample_tasks.add(ev, ev.sample_with_count.remote()) + + with self.timers["replay_processing"]: + for ra, replay in self.replay_tasks.completed(): + self.replay_tasks.add(ra, ra.replay.remote()) + if self.learner.inqueue.full(): + self.num_samples_dropped += 1 + else: + with self.timers["get_samples"]: + samples = ray.get(replay) + # Defensive copy against plasma crashes, see #2610 #3452 + self.learner.inqueue.put((ra, samples and samples.copy())) + + with self.timers["update_priorities"]: + while not self.learner.outqueue.empty(): + ra, prio_dict, count = self.learner.outqueue.get() + ra.update_priorities.remote(prio_dict) + train_timesteps += count + + return sample_timesteps, train_timesteps + + @ray.remote(num_cpus=0) class ReplayActor(object): """A replay buffer shard. @@ -157,182 +341,3 @@ class LearnerThread(threading.Thread): self.outqueue.put((ra, prio_dict, replay.count)) self.learner_queue_size.push(self.inqueue.qsize()) self.weights_updated = True - - -class AsyncReplayOptimizer(PolicyOptimizer): - """Main event loop of the Ape-X optimizer (async sampling with replay). - - This class coordinates the data transfers between the learner thread, - remote evaluators (Ape-X actors), and replay buffer actors. - - This optimizer requires that policy evaluators return an additional - "td_error" array in the info return of compute_gradients(). This error - term will be used for sample prioritization.""" - - def _init(self, - learning_starts=1000, - buffer_size=10000, - prioritized_replay=True, - prioritized_replay_alpha=0.6, - prioritized_replay_beta=0.4, - prioritized_replay_eps=1e-6, - train_batch_size=512, - sample_batch_size=50, - num_replay_buffer_shards=1, - max_weight_sync_delay=400, - debug=False): - - self.debug = debug - self.replay_starts = learning_starts - self.prioritized_replay_beta = prioritized_replay_beta - self.prioritized_replay_eps = prioritized_replay_eps - self.max_weight_sync_delay = max_weight_sync_delay - - self.learner = LearnerThread(self.local_evaluator) - self.learner.start() - - self.replay_actors = create_colocated(ReplayActor, [ - num_replay_buffer_shards, - learning_starts, - buffer_size, - train_batch_size, - prioritized_replay_alpha, - prioritized_replay_beta, - prioritized_replay_eps, - ], num_replay_buffer_shards) - - # Stats - self.timers = { - k: TimerStat() - for k in [ - "put_weights", "get_samples", "sample_processing", - "replay_processing", "update_priorities", "train", "sample" - ] - } - self.num_weight_syncs = 0 - self.num_samples_dropped = 0 - self.learning_started = False - - # Number of worker steps since the last weight update - self.steps_since_update = {} - - # Otherwise kick of replay tasks for local gradient updates - self.replay_tasks = TaskPool() - for ra in self.replay_actors: - for _ in range(REPLAY_QUEUE_DEPTH): - self.replay_tasks.add(ra, ra.replay.remote()) - - # Kick off async background sampling - self.sample_tasks = TaskPool() - if self.remote_evaluators: - self.set_evaluators(self.remote_evaluators) - - # For https://github.com/ray-project/ray/issues/2541 only - def set_evaluators(self, remote_evaluators): - self.remote_evaluators = remote_evaluators - weights = self.local_evaluator.get_weights() - for ev in self.remote_evaluators: - ev.set_weights.remote(weights) - self.steps_since_update[ev] = 0 - for _ in range(SAMPLE_QUEUE_DEPTH): - self.sample_tasks.add(ev, ev.sample_with_count.remote()) - - def step(self): - assert len(self.remote_evaluators) > 0 - start = time.time() - sample_timesteps, train_timesteps = self._step() - time_delta = time.time() - start - self.timers["sample"].push(time_delta) - self.timers["sample"].push_units_processed(sample_timesteps) - if train_timesteps > 0: - self.learning_started = True - if self.learning_started: - self.timers["train"].push(time_delta) - self.timers["train"].push_units_processed(train_timesteps) - self.num_steps_sampled += sample_timesteps - self.num_steps_trained += train_timesteps - - def _step(self): - sample_timesteps, train_timesteps = 0, 0 - weights = None - - with self.timers["sample_processing"]: - completed = list(self.sample_tasks.completed()) - counts = ray.get([c[1][1] for c in completed]) - for i, (ev, (sample_batch, count)) in enumerate(completed): - sample_timesteps += counts[i] - - # Send the data to the replay buffer - random.choice( - self.replay_actors).add_batch.remote(sample_batch) - - # Update weights if needed - self.steps_since_update[ev] += counts[i] - if self.steps_since_update[ev] >= self.max_weight_sync_delay: - # Note that it's important to pull new weights once - # updated to avoid excessive correlation between actors - if weights is None or self.learner.weights_updated: - self.learner.weights_updated = False - with self.timers["put_weights"]: - weights = ray.put( - self.local_evaluator.get_weights()) - ev.set_weights.remote(weights) - self.num_weight_syncs += 1 - self.steps_since_update[ev] = 0 - - # Kick off another sample request - self.sample_tasks.add(ev, ev.sample_with_count.remote()) - - with self.timers["replay_processing"]: - for ra, replay in self.replay_tasks.completed(): - self.replay_tasks.add(ra, ra.replay.remote()) - if self.learner.inqueue.full(): - self.num_samples_dropped += 1 - else: - with self.timers["get_samples"]: - samples = ray.get(replay) - # Defensive copy against plasma crashes, see #2610 #3452 - self.learner.inqueue.put((ra, samples and samples.copy())) - - with self.timers["update_priorities"]: - while not self.learner.outqueue.empty(): - ra, prio_dict, count = self.learner.outqueue.get() - ra.update_priorities.remote(prio_dict) - train_timesteps += count - - return sample_timesteps, train_timesteps - - def stop(self): - for r in self.replay_actors: - r.__ray_terminate__.remote() - self.learner.stopped = True - - def stats(self): - replay_stats = ray.get(self.replay_actors[0].stats.remote(self.debug)) - timing = { - "{}_time_ms".format(k): round(1000 * self.timers[k].mean, 3) - for k in self.timers - } - timing["learner_grad_time_ms"] = round( - 1000 * self.learner.grad_timer.mean, 3) - timing["learner_dequeue_time_ms"] = round( - 1000 * self.learner.queue_timer.mean, 3) - stats = { - "sample_throughput": round(self.timers["sample"].mean_throughput, - 3), - "train_throughput": round(self.timers["train"].mean_throughput, 3), - "num_weight_syncs": self.num_weight_syncs, - "num_samples_dropped": self.num_samples_dropped, - "learner_queue": self.learner.learner_queue_size.stats(), - "replay_shard_0": replay_stats, - } - debug_stats = { - "timing_breakdown": timing, - "pending_sample_tasks": self.sample_tasks.count, - "pending_replay_tasks": self.replay_tasks.count, - } - if self.debug: - stats.update(debug_stats) - if self.learner.stats: - stats["learner"] = self.learner.stats - return dict(PolicyOptimizer.stats(self), **stats) diff --git a/python/ray/rllib/optimizers/async_samples_optimizer.py b/python/ray/rllib/optimizers/async_samples_optimizer.py index 6b8f6014d..ad0d86dfc 100644 --- a/python/ray/rllib/optimizers/async_samples_optimizer.py +++ b/python/ray/rllib/optimizers/async_samples_optimizer.py @@ -18,6 +18,7 @@ import ray from ray.rllib.optimizers.multi_gpu_impl import LocalSyncParallelOptimizer from ray.rllib.optimizers.policy_optimizer import PolicyOptimizer from ray.rllib.utils.actors import TaskPool +from ray.rllib.utils.annotations import override from ray.rllib.utils.timer import TimerStat from ray.rllib.utils.window_stat import WindowStat @@ -27,6 +28,189 @@ LEARNER_QUEUE_MAX_SIZE = 16 NUM_DATA_LOAD_THREADS = 16 +class AsyncSamplesOptimizer(PolicyOptimizer): + """Main event loop of the IMPALA architecture. + + This class coordinates the data transfers between the learner thread + and remote evaluators (IMPALA actors). + """ + + @override(PolicyOptimizer) + def _init(self, + train_batch_size=500, + sample_batch_size=50, + num_envs_per_worker=1, + num_gpus=0, + lr=0.0005, + grad_clip=40, + replay_buffer_num_slots=0, + replay_proportion=0.0, + num_parallel_data_loaders=1, + max_sample_requests_in_flight_per_worker=2, + broadcast_interval=1): + self.learning_started = False + self.train_batch_size = train_batch_size + self.sample_batch_size = sample_batch_size + self.broadcast_interval = broadcast_interval + + if num_gpus > 1 or num_parallel_data_loaders > 1: + logger.info( + "Enabling multi-GPU mode, {} GPUs, {} parallel loaders".format( + num_gpus, num_parallel_data_loaders)) + if train_batch_size // max(1, num_gpus) % ( + sample_batch_size // num_envs_per_worker) != 0: + raise ValueError( + "Sample batches must evenly divide across GPUs.") + self.learner = TFMultiGPULearner( + self.local_evaluator, + lr=lr, + num_gpus=num_gpus, + train_batch_size=train_batch_size, + grad_clip=grad_clip, + num_parallel_data_loaders=num_parallel_data_loaders) + else: + self.learner = LearnerThread(self.local_evaluator) + self.learner.start() + + assert len(self.remote_evaluators) > 0 + + # Stats + self.timers = {k: TimerStat() for k in ["train", "sample"]} + self.num_weight_syncs = 0 + self.num_replayed = 0 + self.learning_started = False + + # Kick off async background sampling + self.sample_tasks = TaskPool() + weights = self.local_evaluator.get_weights() + for ev in self.remote_evaluators: + ev.set_weights.remote(weights) + for _ in range(max_sample_requests_in_flight_per_worker): + self.sample_tasks.add(ev, ev.sample.remote()) + + self.batch_buffer = [] + + if replay_proportion: + assert replay_buffer_num_slots > 0 + assert (replay_buffer_num_slots * sample_batch_size > + train_batch_size) + self.replay_proportion = replay_proportion + self.replay_buffer_num_slots = replay_buffer_num_slots + self.replay_batches = [] + + @override(PolicyOptimizer) + def step(self): + assert self.learner.is_alive() + start = time.time() + sample_timesteps, train_timesteps = self._step() + time_delta = time.time() - start + self.timers["sample"].push(time_delta) + self.timers["sample"].push_units_processed(sample_timesteps) + if train_timesteps > 0: + self.learning_started = True + if self.learning_started: + self.timers["train"].push(time_delta) + self.timers["train"].push_units_processed(train_timesteps) + self.num_steps_sampled += sample_timesteps + self.num_steps_trained += train_timesteps + + @override(PolicyOptimizer) + def stop(self): + self.learner.stopped = True + + @override(PolicyOptimizer) + def stats(self): + timing = { + "{}_time_ms".format(k): round(1000 * self.timers[k].mean, 3) + for k in self.timers + } + timing["learner_grad_time_ms"] = round( + 1000 * self.learner.grad_timer.mean, 3) + timing["learner_load_time_ms"] = round( + 1000 * self.learner.load_timer.mean, 3) + timing["learner_load_wait_time_ms"] = round( + 1000 * self.learner.load_wait_timer.mean, 3) + timing["learner_dequeue_time_ms"] = round( + 1000 * self.learner.queue_timer.mean, 3) + stats = { + "sample_throughput": round(self.timers["sample"].mean_throughput, + 3), + "train_throughput": round(self.timers["train"].mean_throughput, 3), + "num_weight_syncs": self.num_weight_syncs, + "num_steps_replayed": self.num_replayed, + "timing_breakdown": timing, + "learner_queue": self.learner.learner_queue_size.stats(), + } + if self.learner.stats: + stats["learner"] = self.learner.stats + return dict(PolicyOptimizer.stats(self), **stats) + + def _step(self): + sample_timesteps, train_timesteps = 0, 0 + num_sent = 0 + weights = None + + for ev, sample_batch in self._augment_with_replay( + self.sample_tasks.completed_prefetch()): + self.batch_buffer.append(sample_batch) + if sum(b.count + for b in self.batch_buffer) >= self.train_batch_size: + train_batch = self.batch_buffer[0].concat_samples( + self.batch_buffer) + self.learner.inqueue.put(train_batch) + self.batch_buffer = [] + + # If the batch was replayed, skip the update below. + if ev is None: + continue + + sample_timesteps += sample_batch.count + + # Put in replay buffer if enabled + if self.replay_buffer_num_slots > 0: + self.replay_batches.append(sample_batch) + if len(self.replay_batches) > self.replay_buffer_num_slots: + self.replay_batches.pop(0) + + # Note that it's important to pull new weights once + # updated to avoid excessive correlation between actors + if weights is None or (self.learner.weights_updated + and num_sent >= self.broadcast_interval): + self.learner.weights_updated = False + weights = ray.put(self.local_evaluator.get_weights()) + num_sent = 0 + ev.set_weights.remote(weights) + self.num_weight_syncs += 1 + num_sent += 1 + + # Kick off another sample request + self.sample_tasks.add(ev, ev.sample.remote()) + + while not self.learner.outqueue.empty(): + count = self.learner.outqueue.get() + train_timesteps += count + + return sample_timesteps, train_timesteps + + def _augment_with_replay(self, sample_futures): + def can_replay(): + num_needed = int( + np.ceil(self.train_batch_size / self.sample_batch_size)) + return len(self.replay_batches) > num_needed + + for ev, sample_batch in sample_futures: + sample_batch = ray.get(sample_batch) + yield ev, sample_batch + + if can_replay(): + f = self.replay_proportion + while random.random() < f: + f -= 1 + replay_batch = random.choice(self.replay_batches) + self.num_replayed += replay_batch.count + yield None, replay_batch + + class LearnerThread(threading.Thread): """Background thread that updates the local model from sample trajectories. @@ -112,7 +296,7 @@ class TFMultiGPULearner(LearnerThread): LocalSyncParallelOptimizer( adam, self.devices, - [v for _, v in self.policy.loss_inputs()], + [v for _, v in self.policy._loss_inputs], rnn_inputs, 999999, # it will get rounded down self.policy.copy, @@ -129,6 +313,7 @@ class TFMultiGPULearner(LearnerThread): self.loader_thread = _LoaderThread(self, share_stats=(i == 0)) self.loader_thread.start() + @override(LearnerThread) def step(self): assert self.loader_thread.is_alive() with self.load_wait_timer: @@ -158,9 +343,9 @@ class _LoaderThread(threading.Thread): def run(self): while True: - self.step() + self._step() - def step(self): + def _step(self): s = self.learner with self.queue_timer: batch = s.inqueue.get() @@ -169,7 +354,7 @@ class _LoaderThread(threading.Thread): with self.load_timer: tuples = s.policy._get_loss_inputs_dict(batch) - data_keys = [ph for _, ph in s.policy.loss_inputs()] + data_keys = [ph for _, ph in s.policy._loss_inputs] if s.policy._state_inputs: state_keys = s.policy._state_inputs + [s.policy._seq_lens] else: @@ -178,182 +363,3 @@ class _LoaderThread(threading.Thread): [tuples[k] for k in state_keys]) s.ready_optimizers.put(opt) - - -class AsyncSamplesOptimizer(PolicyOptimizer): - """Main event loop of the IMPALA architecture. - - This class coordinates the data transfers between the learner thread - and remote evaluators (IMPALA actors). - """ - - def _init(self, - train_batch_size=500, - sample_batch_size=50, - num_envs_per_worker=1, - num_gpus=0, - lr=0.0005, - grad_clip=40, - replay_buffer_num_slots=0, - replay_proportion=0.0, - num_parallel_data_loaders=1, - max_sample_requests_in_flight_per_worker=2, - broadcast_interval=1): - self.learning_started = False - self.train_batch_size = train_batch_size - self.sample_batch_size = sample_batch_size - self.broadcast_interval = broadcast_interval - - if num_gpus > 1 or num_parallel_data_loaders > 1: - logger.info( - "Enabling multi-GPU mode, {} GPUs, {} parallel loaders".format( - num_gpus, num_parallel_data_loaders)) - if train_batch_size // max(1, num_gpus) % ( - sample_batch_size // num_envs_per_worker) != 0: - raise ValueError( - "Sample batches must evenly divide across GPUs.") - self.learner = TFMultiGPULearner( - self.local_evaluator, - lr=lr, - num_gpus=num_gpus, - train_batch_size=train_batch_size, - grad_clip=grad_clip, - num_parallel_data_loaders=num_parallel_data_loaders) - else: - self.learner = LearnerThread(self.local_evaluator) - self.learner.start() - - assert len(self.remote_evaluators) > 0 - - # Stats - self.timers = {k: TimerStat() for k in ["train", "sample"]} - self.num_weight_syncs = 0 - self.num_replayed = 0 - self.learning_started = False - - # Kick off async background sampling - self.sample_tasks = TaskPool() - weights = self.local_evaluator.get_weights() - for ev in self.remote_evaluators: - ev.set_weights.remote(weights) - for _ in range(max_sample_requests_in_flight_per_worker): - self.sample_tasks.add(ev, ev.sample.remote()) - - self.batch_buffer = [] - - if replay_proportion: - assert replay_buffer_num_slots > 0 - assert (replay_buffer_num_slots * sample_batch_size > - train_batch_size) - self.replay_proportion = replay_proportion - self.replay_buffer_num_slots = replay_buffer_num_slots - self.replay_batches = [] - - def step(self): - assert self.learner.is_alive() - start = time.time() - sample_timesteps, train_timesteps = self._step() - time_delta = time.time() - start - self.timers["sample"].push(time_delta) - self.timers["sample"].push_units_processed(sample_timesteps) - if train_timesteps > 0: - self.learning_started = True - if self.learning_started: - self.timers["train"].push(time_delta) - self.timers["train"].push_units_processed(train_timesteps) - self.num_steps_sampled += sample_timesteps - self.num_steps_trained += train_timesteps - - def _augment_with_replay(self, sample_futures): - def can_replay(): - num_needed = int( - np.ceil(self.train_batch_size / self.sample_batch_size)) - return len(self.replay_batches) > num_needed - - for ev, sample_batch in sample_futures: - sample_batch = ray.get(sample_batch) - yield ev, sample_batch - - if can_replay(): - f = self.replay_proportion - while random.random() < f: - f -= 1 - replay_batch = random.choice(self.replay_batches) - self.num_replayed += replay_batch.count - yield None, replay_batch - - def _step(self): - sample_timesteps, train_timesteps = 0, 0 - num_sent = 0 - weights = None - - for ev, sample_batch in self._augment_with_replay( - self.sample_tasks.completed_prefetch()): - self.batch_buffer.append(sample_batch) - if sum(b.count - for b in self.batch_buffer) >= self.train_batch_size: - train_batch = self.batch_buffer[0].concat_samples( - self.batch_buffer) - self.learner.inqueue.put(train_batch) - self.batch_buffer = [] - - # If the batch was replayed, skip the update below. - if ev is None: - continue - - sample_timesteps += sample_batch.count - - # Put in replay buffer if enabled - if self.replay_buffer_num_slots > 0: - self.replay_batches.append(sample_batch) - if len(self.replay_batches) > self.replay_buffer_num_slots: - self.replay_batches.pop(0) - - # Note that it's important to pull new weights once - # updated to avoid excessive correlation between actors - if weights is None or (self.learner.weights_updated - and num_sent >= self.broadcast_interval): - self.learner.weights_updated = False - weights = ray.put(self.local_evaluator.get_weights()) - num_sent = 0 - ev.set_weights.remote(weights) - self.num_weight_syncs += 1 - num_sent += 1 - - # Kick off another sample request - self.sample_tasks.add(ev, ev.sample.remote()) - - while not self.learner.outqueue.empty(): - count = self.learner.outqueue.get() - train_timesteps += count - - return sample_timesteps, train_timesteps - - def stop(self): - self.learner.stopped = True - - def stats(self): - timing = { - "{}_time_ms".format(k): round(1000 * self.timers[k].mean, 3) - for k in self.timers - } - timing["learner_grad_time_ms"] = round( - 1000 * self.learner.grad_timer.mean, 3) - timing["learner_load_time_ms"] = round( - 1000 * self.learner.load_timer.mean, 3) - timing["learner_load_wait_time_ms"] = round( - 1000 * self.learner.load_wait_timer.mean, 3) - timing["learner_dequeue_time_ms"] = round( - 1000 * self.learner.queue_timer.mean, 3) - stats = { - "sample_throughput": round(self.timers["sample"].mean_throughput, - 3), - "train_throughput": round(self.timers["train"].mean_throughput, 3), - "num_weight_syncs": self.num_weight_syncs, - "num_steps_replayed": self.num_replayed, - "timing_breakdown": timing, - "learner_queue": self.learner.learner_queue_size.stats(), - } - if self.learner.stats: - stats["learner"] = self.learner.stats - return dict(PolicyOptimizer.stats(self), **stats) diff --git a/python/ray/rllib/optimizers/multi_gpu_optimizer.py b/python/ray/rllib/optimizers/multi_gpu_optimizer.py index 771acb5ac..a3ec1ff92 100644 --- a/python/ray/rllib/optimizers/multi_gpu_optimizer.py +++ b/python/ray/rllib/optimizers/multi_gpu_optimizer.py @@ -12,6 +12,7 @@ import ray from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph from ray.rllib.optimizers.policy_optimizer import PolicyOptimizer from ray.rllib.optimizers.multi_gpu_impl import LocalSyncParallelOptimizer +from ray.rllib.utils.annotations import override from ray.rllib.utils.timer import TimerStat logger = logging.getLogger(__name__) @@ -33,6 +34,7 @@ class LocalMultiGPUOptimizer(PolicyOptimizer): may result in unexpected behavior. """ + @override(PolicyOptimizer) def _init(self, sgd_batch_size=128, num_sgd_iter=10, @@ -85,12 +87,13 @@ class LocalMultiGPUOptimizer(PolicyOptimizer): rnn_inputs = [] self.par_opt = LocalSyncParallelOptimizer( self.policy.optimizer(), self.devices, - [v for _, v in self.policy.loss_inputs()], rnn_inputs, + [v for _, v in self.policy._loss_inputs], rnn_inputs, self.per_device_batch_size, self.policy.copy) self.sess = self.local_evaluator.tf_sess self.sess.run(tf.global_variables_initializer()) + @override(PolicyOptimizer) def step(self): with self.update_weights_timer: if self.remote_evaluators: @@ -119,7 +122,7 @@ class LocalMultiGPUOptimizer(PolicyOptimizer): with self.load_timer: tuples = self.policy._get_loss_inputs_dict(samples) - data_keys = [ph for _, ph in self.policy.loss_inputs()] + data_keys = [ph for _, ph in self.policy._loss_inputs] if self.policy._state_inputs: state_keys = ( self.policy._state_inputs + [self.policy._seq_lens]) @@ -148,6 +151,7 @@ class LocalMultiGPUOptimizer(PolicyOptimizer): self.num_steps_trained += samples.count return _averaged(iter_extra_fetches) + @override(PolicyOptimizer) def stats(self): return dict( PolicyOptimizer.stats(self), **{ diff --git a/python/ray/rllib/optimizers/policy_optimizer.py b/python/ray/rllib/optimizers/policy_optimizer.py index 6c2a5787d..adac0258f 100644 --- a/python/ray/rllib/optimizers/policy_optimizer.py +++ b/python/ray/rllib/optimizers/policy_optimizer.py @@ -63,7 +63,7 @@ class PolicyOptimizer(object): def _init(self): """Subclasses should prefer overriding this instead of __init__.""" - pass + raise NotImplementedError def step(self): """Takes a logical optimization step. @@ -86,6 +86,21 @@ class PolicyOptimizer(object): "num_steps_sampled": self.num_steps_sampled, } + def save(self): + """Returns a serializable object representing the optimizer state.""" + + return [self.num_steps_trained, self.num_steps_sampled] + + def restore(self, data): + """Restores optimizer state from the given data object.""" + + self.num_steps_trained = data[0] + self.num_steps_sampled = data[1] + + def stop(self): + """Release any resources used by this optimizer.""" + pass + def collect_metrics(self, timeout_seconds, min_history=100, @@ -118,17 +133,6 @@ class PolicyOptimizer(object): res.update(info=self.stats()) return res - def save(self): - """Returns a serializable object representing the optimizer state.""" - - return [self.num_steps_trained, self.num_steps_sampled] - - def restore(self, data): - """Restores optimizer state from the given data object.""" - - self.num_steps_trained = data[0] - self.num_steps_sampled = data[1] - def foreach_evaluator(self, func): """Apply the given function to each evaluator instance.""" @@ -150,10 +154,6 @@ class PolicyOptimizer(object): ]) return local_result + remote_results - def stop(self): - """Release any resources used by this optimizer.""" - pass - @staticmethod def _check_not_multiagent(sample_batch): if isinstance(sample_batch, MultiAgentBatch): diff --git a/python/ray/rllib/optimizers/sync_replay_optimizer.py b/python/ray/rllib/optimizers/sync_replay_optimizer.py index b15ab3390..f2a42a083 100644 --- a/python/ray/rllib/optimizers/sync_replay_optimizer.py +++ b/python/ray/rllib/optimizers/sync_replay_optimizer.py @@ -11,6 +11,7 @@ from ray.rllib.optimizers.replay_buffer import ReplayBuffer, \ from ray.rllib.optimizers.policy_optimizer import PolicyOptimizer from ray.rllib.evaluation.sample_batch import SampleBatch, DEFAULT_POLICY_ID, \ MultiAgentBatch +from ray.rllib.utils.annotations import override from ray.rllib.utils.compression import pack_if_needed from ray.rllib.utils.filter import RunningStat from ray.rllib.utils.timer import TimerStat @@ -24,6 +25,7 @@ class SyncReplayOptimizer(PolicyOptimizer): "td_error" array in the info return of compute_gradients(). This error term will be used for sample prioritization.""" + @override(PolicyOptimizer) def _init(self, learning_starts=1000, buffer_size=10000, @@ -70,6 +72,7 @@ class SyncReplayOptimizer(PolicyOptimizer): assert buffer_size >= self.replay_starts + @override(PolicyOptimizer) def step(self): with self.update_weights_timer: if self.remote_evaluators: @@ -106,6 +109,21 @@ class SyncReplayOptimizer(PolicyOptimizer): self.num_steps_sampled += batch.count + @override(PolicyOptimizer) + def stats(self): + return dict( + PolicyOptimizer.stats(self), **{ + "sample_time_ms": round(1000 * self.sample_timer.mean, 3), + "replay_time_ms": round(1000 * self.replay_timer.mean, 3), + "grad_time_ms": round(1000 * self.grad_timer.mean, 3), + "update_time_ms": round(1000 * self.update_weights_timer.mean, + 3), + "opt_peak_throughput": round(self.grad_timer.mean_throughput, + 3), + "opt_samples": round(self.grad_timer.mean_units_processed, 3), + "learner": self.learner_stats, + }) + def _optimize(self): samples = self._replay() @@ -151,17 +169,3 @@ class SyncReplayOptimizer(PolicyOptimizer): "batch_indexes": batch_indexes }) return MultiAgentBatch(samples, self.train_batch_size) - - def stats(self): - return dict( - PolicyOptimizer.stats(self), **{ - "sample_time_ms": round(1000 * self.sample_timer.mean, 3), - "replay_time_ms": round(1000 * self.replay_timer.mean, 3), - "grad_time_ms": round(1000 * self.grad_timer.mean, 3), - "update_time_ms": round(1000 * self.update_weights_timer.mean, - 3), - "opt_peak_throughput": round(self.grad_timer.mean_throughput, - 3), - "opt_samples": round(self.grad_timer.mean_units_processed, 3), - "learner": self.learner_stats, - }) diff --git a/python/ray/rllib/optimizers/sync_samples_optimizer.py b/python/ray/rllib/optimizers/sync_samples_optimizer.py index 38d5269f0..b78e3ed01 100644 --- a/python/ray/rllib/optimizers/sync_samples_optimizer.py +++ b/python/ray/rllib/optimizers/sync_samples_optimizer.py @@ -6,6 +6,7 @@ import ray import logging from ray.rllib.optimizers.policy_optimizer import PolicyOptimizer from ray.rllib.evaluation.sample_batch import SampleBatch +from ray.rllib.utils.annotations import override from ray.rllib.utils.filter import RunningStat from ray.rllib.utils.timer import TimerStat @@ -20,6 +21,7 @@ class SyncSamplesOptimizer(PolicyOptimizer): model weights are then broadcast to all remote evaluators. """ + @override(PolicyOptimizer) def _init(self, num_sgd_iter=1, train_batch_size=1): self.update_weights_timer = TimerStat() self.sample_timer = TimerStat() @@ -29,6 +31,7 @@ class SyncSamplesOptimizer(PolicyOptimizer): self.train_batch_size = train_batch_size self.learner_stats = {} + @override(PolicyOptimizer) def step(self): with self.update_weights_timer: if self.remote_evaluators: @@ -62,6 +65,7 @@ class SyncSamplesOptimizer(PolicyOptimizer): self.num_steps_trained += samples.count return fetches + @override(PolicyOptimizer) def stats(self): return dict( PolicyOptimizer.stats(self), **{ diff --git a/python/ray/rllib/test/test_evaluators.py b/python/ray/rllib/test/test_evaluators.py index 9ae0994f3..c7a72d7a5 100644 --- a/python/ray/rllib/test/test_evaluators.py +++ b/python/ray/rllib/test/test_evaluators.py @@ -4,7 +4,7 @@ from __future__ import print_function import unittest -from ray.rllib.agents.dqn.dqn_policy_graph import adjust_nstep +from ray.rllib.agents.dqn.dqn_policy_graph import _adjust_nstep class DQNTest(unittest.TestCase): @@ -14,7 +14,7 @@ class DQNTest(unittest.TestCase): rewards = [10.0, 0.0, 100.0, 100.0, 100.0, 100.0, 100.0] new_obs = [2, 3, 4, 5, 6, 7, 8] dones = [0, 0, 0, 0, 0, 0, 1] - adjust_nstep(3, 0.9, obs, actions, rewards, new_obs, dones) + _adjust_nstep(3, 0.9, obs, actions, rewards, new_obs, dones) self.assertEqual(obs, [1, 2, 3, 4, 5, 6, 7]) self.assertEqual(actions, ["a", "b", "a", "a", "a", "b", "a"]) self.assertEqual(new_obs, [4, 5, 6, 7, 8, 8, 8]) diff --git a/python/ray/rllib/utils/annotations.py b/python/ray/rllib/utils/annotations.py new file mode 100644 index 000000000..d68f76a69 --- /dev/null +++ b/python/ray/rllib/utils/annotations.py @@ -0,0 +1,20 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + + +def override(cls): + """Annotation for documenting method overrides. + + Arguments: + cls (type): The superclass that provides the overriden method. If this + cls does not actually have the method, an error is raised. + """ + + def check_override(method): + if method.__name__ not in dir(cls): + raise NameError("{} does not override any method of {}".format( + method, cls)) + return method + + return check_override