wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-06 05:16:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

[rllib] Better document which methods are abstract and which ones are overrides (#3480)

Browse Source
This commit is contained in:
Eric Liang
2018-12-08 16:28:58 -08:00
committed by GitHub
parent 462e6ef066
commit 8b5827b9da
40 changed files with 1385 additions and 1167 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/rllib/agents/a3c/a2c.py
+2
View File
@@ -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,
python/ray/rllib/agents/a3c/a3c.py
+8 -5
View File
@@ -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"])
python/ray/rllib/agents/a3c/a3c_tf_policy_graph.py
+29 -22
View File
@@ -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]
python/ray/rllib/agents/a3c/a3c_torch_policy_graph.py
+5
View File
@@ -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,
python/ray/rllib/agents/agent.py
+183 -176
View File
@@ -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):
python/ray/rllib/agents/ars/ars.py
+30 -25
View File
@@ -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]
python/ray/rllib/agents/ddpg/apex.py
+2
View File
@@ -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 > \
python/ray/rllib/agents/ddpg/ddpg.py
+2
View File
@@ -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"]:
python/ray/rllib/agents/ddpg/ddpg_policy_graph.py
+71 -60
View File
@@ -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])
python/ray/rllib/agents/dqn/apex.py
+2
View File
@@ -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 > \
python/ray/rllib/agents/dqn/dqn.py
+36 -33
View File
@@ -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({
python/ray/rllib/agents/dqn/dqn_policy_graph.py
+76 -67
View File
@@ -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,
python/ray/rllib/agents/es/es.py
+30 -25
View File
@@ -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]
python/ray/rllib/agents/impala/impala.py
+3
View File
@@ -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()
python/ray/rllib/agents/impala/vtrace_policy_graph.py
+16 -7
View File
@@ -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)
python/ray/rllib/agents/pg/pg.py
+3
View File
@@ -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()
python/ray/rllib/agents/pg/pg_policy_graph.py
+4
View File
@@ -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
python/ray/rllib/agents/ppo/ppo.py
+21 -18
View File
@@ -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
python/ray/rllib/agents/ppo/ppo_policy_graph.py
+32 -24
View File
@@ -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]
python/ray/rllib/env/async_vector_env.py
Vendored
+15 -5
View File
@@ -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"
python/ray/rllib/env/vector_env.py
Vendored
+6
View File
@@ -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
python/ray/rllib/evaluation/policy_evaluator.py
+109 -102
View File
@@ -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):
python/ray/rllib/evaluation/sampler.py
+1 -1
View File
@@ -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],
python/ray/rllib/evaluation/tf_policy_graph.py
+125 -105
View File
@@ -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)
python/ray/rllib/evaluation/torch_policy_graph.py
+17 -11
View File
@@ -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())
python/ray/rllib/models/action_dist.py
+15 -6
View File
@@ -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])
python/ray/rllib/models/fcnet.py
+2
View File
@@ -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.
python/ray/rllib/models/lstm.py
+66 -64
View File
@@ -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
python/ray/rllib/models/model.py
+13 -13
View File
@@ -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"):
python/ray/rllib/models/preprocessors.py
+14
View File
@@ -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())))
python/ray/rllib/models/visionnet.py
+4 -2
View File
@@ -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],
python/ray/rllib/optimizers/async_gradients_optimizer.py
+4
View File
@@ -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), **{
python/ray/rllib/optimizers/async_replay_optimizer.py
+184 -179
View File
@@ -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)
python/ray/rllib/optimizers/async_samples_optimizer.py
+189 -183
View File
@@ -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)
python/ray/rllib/optimizers/multi_gpu_optimizer.py
+6 -2
View File
@@ -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), **{
python/ray/rllib/optimizers/policy_optimizer.py
+16 -16
View File
@@ -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):
python/ray/rllib/optimizers/sync_replay_optimizer.py
+18 -14
View File
@@ -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,
})
python/ray/rllib/optimizers/sync_samples_optimizer.py
+4
View File
@@ -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), **{
python/ray/rllib/test/test_evaluators.py
+2 -2
View File
@@ -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])
python/ray/rllib/utils/annotations.py
+20
View File
@@ -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