wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-19 11:27:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

[RLlib] Policy-classes cleanup and torch/tf unification. (#6770)

Browse Source
This commit is contained in:
Sven Mika
2020-01-17 22:26:28 -08:00
committed by Eric Liang
parent 763818b476
commit 303547f119
23 changed files with 282 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/agents/a3c/a3c_torch_policy.py
+2 -1
View File
@@ -48,7 +48,8 @@ def add_advantages(policy,
policy.config["lambda"])
def model_value_predictions(policy, input_dict, state_batches, model):
def model_value_predictions(policy, input_dict, state_batches, model,
action_dist):
return {SampleBatch.VF_PREDS: model.value_function().cpu().numpy()}
rllib/agents/ddpg/ddpg_policy.py
+1
View File
@@ -288,6 +288,7 @@ class DDPGTFPolicy(DDPGPostprocessing, TFPolicy):
self,
observation_space,
action_space,
self.config,
self.sess,
obs_input=self.cur_observations,
action_sampler=self.output_actions,
rllib/agents/dqn/apex.py
+1 -1
View File
@@ -33,7 +33,7 @@ APEX_DEFAULT_CONFIG = merge_dicts(
def defer_make_workers(trainer, env_creator, policy, config):
# Hack to workaround https://github.com/ray-project/ray/issues/2541
# The workers will be creatd later, after the optimizer is created
# The workers will be created later, after the optimizer is created
return trainer._make_workers(env_creator, policy, config, 0)
rllib/agents/marwil/marwil_policy.py
+1
View File
@@ -134,6 +134,7 @@ class MARWILPolicy(MARWILPostprocessing, TFPolicy):
self,
observation_space,
action_space,
self.config,
self.sess,
obs_input=self.obs_t,
action_sampler=self.output_actions,
rllib/agents/ppo/ppo.py
+2 -1
View File
@@ -138,7 +138,8 @@ def warn_about_bad_reward_scales(trainer, result):
"This means that it will take more than "
"{} iterations for your value ".format(rew_scale) +
"function to converge. If this is not intended, consider "
"increasing `vf_clip_param`.")
"increasing `vf_clip_param`."
)
def validate_config(config):
rllib/agents/qmix/qmix_policy.py
+2 -4
View File
@@ -143,6 +143,7 @@ class QMixLoss(nn.Module):
return loss, mask, masked_td_error, chosen_action_qvals, targets
# TODO(sven): Make this a TorchPolicy child.
class QMixTorchPolicy(Policy):
"""QMix impl. Assumes homogeneous agents for now.
@@ -154,13 +155,10 @@ class QMixTorchPolicy(Policy):
dict space with an action_mask key, e.g. {"obs": ob, "action_mask": mask}.
The mask space must be `Box(0, 1, (n_actions,))`.
"""
def __init__(self, obs_space, action_space, config):
_validate(obs_space, action_space)
config = dict(ray.rllib.agents.qmix.qmix.DEFAULT_CONFIG, **config)
self.config = config
self.observation_space = obs_space
self.action_space = action_space
super().__init__(obs_space, action_space, config)
self.n_agents = len(obs_space.original_space.spaces)
self.n_actions = action_space.spaces[0].n
self.h_size = config["model"]["lstm_cell_size"]
rllib/contrib/alpha_zero/core/alpha_zero_policy.py
+5 -3
View File
@@ -10,11 +10,13 @@ torch, _ = try_import_torch()
class AlphaZeroPolicy(TorchPolicy):
def __init__(self, observation_space, action_space, model, loss,
def __init__(self, observation_space, action_space, config, model, loss,
action_distribution_class, mcts_creator, env_creator,
**kwargs):
super().__init__(observation_space, action_space, model, loss,
action_distribution_class)
super().__init__(
observation_space, action_space, config, model, loss,
action_distribution_class
)
# we maintain an env copy in the policy that is used during mcts
# simulations
self.env_creator = env_creator
rllib/contrib/alpha_zero/core/alpha_zero_trainer.py
+4 -2
View File
@@ -160,8 +160,10 @@ class AlphaZeroPolicyWrapperClass(AlphaZeroPolicy):
def mcts_creator():
return MCTS(model, config["mcts_config"])
super().__init__(obs_space, action_space, model, alpha_zero_loss,
TorchCategorical, mcts_creator, _env_creator)
super().__init__(
obs_space, action_space, config, model, alpha_zero_loss,
TorchCategorical, mcts_creator, _env_creator
)
AlphaZeroTrainer = build_trainer(
rllib/contrib/maddpg/maddpg_policy.py
+28 -17
View File
@@ -47,8 +47,7 @@ class MADDPGPostprocessing:
class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
def __init__(self, obs_space, act_space, config):
# _____ Initial Configuration
self.config = config = dict(ray.rllib.contrib.maddpg.DEFAULT_CONFIG,
**config)
config = dict(ray.rllib.contrib.maddpg.DEFAULT_CONFIG, **config)
self.global_step = tf.train.get_or_create_global_step()
# FIXME: Get done from info is required since agentwise done is not
@@ -120,8 +119,9 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
act_ph_n,
obs_space_n,
act_space_n,
hiddens=config["critic_hiddens"],
activation=getattr(tf.nn, config["critic_hidden_activation"]),
config["use_state_preprocessor"],
config["critic_hiddens"],
getattr(tf.nn, config["critic_hidden_activation"]),
scope="critic")
# Build critic network for t + 1.
@@ -130,8 +130,9 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
new_act_ph_n,
obs_space_n,
act_space_n,
hiddens=config["critic_hiddens"],
activation=getattr(tf.nn, config["critic_hidden_activation"]),
config["use_state_preprocessor"],
config["critic_hiddens"],
getattr(tf.nn, config["critic_hidden_activation"]),
scope="target_critic")
# Build critic loss.
@@ -149,8 +150,9 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
obs_ph_n[agent_id],
obs_space_n[agent_id],
act_space_n[agent_id],
hiddens=config["actor_hiddens"],
activation=getattr(tf.nn, config["actor_hidden_activation"]),
config["use_state_preprocessor"],
config["actor_hiddens"],
getattr(tf.nn, config["actor_hidden_activation"]),
scope="actor"))
# Build actor network for t + 1.
@@ -160,8 +162,9 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
self.new_obs_ph,
obs_space_n[agent_id],
act_space_n[agent_id],
hiddens=config["actor_hiddens"],
activation=getattr(tf.nn, config["actor_hidden_activation"]),
config["use_state_preprocessor"],
config["actor_hiddens"],
getattr(tf.nn, config["actor_hidden_activation"]),
scope="target_actor"))
# Build actor loss.
@@ -172,8 +175,9 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
act_n,
obs_space_n,
act_space_n,
hiddens=config["critic_hiddens"],
activation=getattr(tf.nn, config["critic_hidden_activation"]),
config["use_state_preprocessor"],
config["critic_hiddens"],
getattr(tf.nn, config["critic_hidden_activation"]),
scope="critic")
actor_loss = -tf.reduce_mean(critic)
if config["actor_feature_reg"] is not None:
@@ -238,7 +242,8 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
self,
obs_space,
act_space,
self.sess,
config=config,
sess=self.sess,
obs_input=obs_ph_n[agent_id],
action_sampler=act_sampler,
loss=actor_loss + critic_loss,
@@ -313,11 +318,12 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
act_n,
obs_space_n,
act_space_n,
use_state_preprocessor,
hiddens,
activation=None,
scope=None):
with tf.variable_scope(scope, reuse=tf.AUTO_REUSE) as scope:
if self.config["use_state_preprocessor"]:
if use_state_preprocessor:
model_n = [
ModelCatalog.get_model({
"obs": obs,
@@ -333,7 +339,9 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
out = tf.concat(obs_n + act_n, axis=1)
for hidden in hiddens:
out = tf.layers.dense(out, units=hidden, activation=activation)
out = tf.layers.dense(
out, units=hidden, activation=activation
)
feature = out
out = tf.layers.dense(feature, units=1, activation=None)
@@ -343,11 +351,12 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
obs,
obs_space,
act_space,
use_state_preprocessor,
hiddens,
activation=None,
scope=None):
with tf.variable_scope(scope, reuse=tf.AUTO_REUSE) as scope:
if self.config["use_state_preprocessor"]:
if use_state_preprocessor:
model = ModelCatalog.get_model({
"obs": obs,
"is_training": self._get_is_training_placeholder(),
@@ -358,7 +367,9 @@ class MADDPGTFPolicy(MADDPGPostprocessing, TFPolicy):
out = obs
for hidden in hiddens:
out = tf.layers.dense(out, units=hidden, activation=activation)
out = tf.layers.dense(
out, units=hidden, activation=activation
)
feature = tf.layers.dense(
out, units=act_space.shape[0], activation=None)
sampler = tfp.distributions.RelaxedOneHotCategorical(
rllib/examples/multiagent_custom_policy.py
+1 -1
View File
@@ -31,7 +31,7 @@ class RandomPolicy(Policy):
def compute_actions(self,
obs_batch,
state_batches,
state_batches=None,
prev_action_batch=None,
prev_reward_batch=None,
info_batch=None,
rllib/examples/rock_paper_scissors_multiagent.py
+2 -9
View File
@@ -78,15 +78,12 @@ class RockPaperScissorsEnv(MultiAgentEnv):
class AlwaysSameHeuristic(Policy):
"""Pick a random move and stick with it for the entire episode."""
def __init__(self, observation_space, action_space, config):
Policy.__init__(self, observation_space, action_space, config)
def get_initial_state(self):
return [random.choice([ROCK, PAPER, SCISSORS])]
def compute_actions(self,
obs_batch,
state_batches,
state_batches=None,
prev_action_batch=None,
prev_reward_batch=None,
info_batch=None,
@@ -106,13 +103,9 @@ class AlwaysSameHeuristic(Policy):
class BeatLastHeuristic(Policy):
"""Play the move that would beat the last move of the opponent."""
def __init__(self, observation_space, action_space, config):
Policy.__init__(self, observation_space, action_space, config)
def compute_actions(self,
obs_batch,
state_batches,
state_batches=None,
prev_action_batch=None,
prev_reward_batch=None,
info_batch=None,
rllib/examples/rollout_worker_custom_workflow.py
+1 -1
View File
@@ -34,7 +34,7 @@ class CustomPolicy(Policy):
def compute_actions(self,
obs_batch,
state_batches,
state_batches=None,
prev_action_batch=None,
prev_reward_batch=None,
info_batch=None,
rllib/policy/dynamic_tf_policy.py
+8 -11
View File
@@ -64,8 +64,8 @@ class DynamicTFPolicy(TFPolicy):
TF fetches given the policy and batch input tensors
grad_stats_fn (func): optional function that returns a dict of
TF fetches given the policy and loss gradient tensors
before_loss_init (func): optional function to run prior to loss
init that takes the same arguments as __init__
before_loss_init (Optional[callable]): Optional function to run
prior to loss init that takes the same arguments as __init__.
make_model (func): optional function that returns a ModelV2 object
given (policy, obs_space, action_space, config).
All policy variables should be created in this function. If not
@@ -74,7 +74,7 @@ class DynamicTFPolicy(TFPolicy):
tuple of action and action logp tensors given
(policy, model, input_dict, obs_space, action_space, config).
If not specified, a default action distribution will be used.
existing_inputs (OrderedDict): when copying a policy, this
existing_inputs (OrderedDict): When copying a policy, this
specifies an existing dict of placeholders to use instead of
defining new ones
existing_model (ModelV2): when copying a policy, this specifies
@@ -176,6 +176,7 @@ class DynamicTFPolicy(TFPolicy):
self,
obs_space,
action_space,
config,
sess,
obs_input=obs,
action_sampler=action_sampler,
@@ -191,8 +192,10 @@ class DynamicTFPolicy(TFPolicy):
max_seq_len=config["model"]["max_seq_len"],
batch_divisibility_req=batch_divisibility_req)
# Phase 2 init
before_loss_init(self, obs_space, action_space, config)
# Phase 2 init.
if before_loss_init is not None:
before_loss_init(self, obs_space, action_space, config)
if not existing_inputs:
self._initialize_loss()
@@ -248,12 +251,6 @@ class DynamicTFPolicy(TFPolicy):
else:
return []
def is_recurrent(self):
return len(self._state_in) > 0
def num_state_tensors(self):
return len(self._state_in)
def _initialize_loss(self):
def fake_array(tensor):
shape = tensor.shape.as_list()
rllib/policy/eager_tf_policy.py
+10 -1
View File
@@ -10,7 +10,7 @@ from ray.rllib.evaluation.episode import _flatten_action
from ray.rllib.models.catalog import ModelCatalog
from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY
from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.policy.tf_policy import ACTION_PROB, ACTION_LOGP
from ray.rllib.policy.policy import ACTION_PROB, ACTION_LOGP
from ray.rllib.utils import add_mixins
from ray.rllib.utils.annotations import override
from ray.rllib.utils.debug import log_once
@@ -365,6 +365,7 @@ def build_eager_tf_policy(name,
def is_recurrent(self):
return len(self._state_in) > 0
@override(Policy)
def num_state_tensors(self):
return len(self._state_in)
@@ -380,6 +381,14 @@ def build_eager_tf_policy(name,
def loss_initialized(self):
return self._loss_initialized
@override(Policy)
def export_model(self, export_dir):
pass
@override(Policy)
def export_checkpoint(self, export_dir):
pass
def _get_is_training_placeholder(self):
return tf.convert_to_tensor(self._is_training)
rllib/policy/policy.py
+43 -23
View File
@@ -1,6 +1,7 @@
from abc import ABCMeta, abstractmethod
from collections import namedtuple
import numpy as np
import gym
import numpy as np
from ray.rllib.utils.annotations import DeveloperAPI
@@ -8,6 +9,9 @@ from ray.rllib.utils.annotations import DeveloperAPI
# `grad_info` dict returned by learn_on_batch() / compute_grads() via this key.
LEARNER_STATS_KEY = "learner_stats"
ACTION_PROB = "action_prob"
ACTION_LOGP = "action_logp"
class TupleActions(namedtuple("TupleActions", ["batches"])):
"""Used to return tuple actions as a list of batches per tuple element."""
@@ -20,7 +24,7 @@ class TupleActions(namedtuple("TupleActions", ["batches"])):
@DeveloperAPI
class Policy:
class Policy(metaclass=ABCMeta):
"""An agent policy and loss, i.e., a TFPolicy or other subclass.
This object defines how to act in the environment, and also losses used to
@@ -51,27 +55,31 @@ class Policy:
action_space (gym.Space): Action space of the policy.
config (dict): Policy-specific configuration data.
"""
self.observation_space = observation_space
self.action_space = action_space
self.config = config
@abstractmethod
@DeveloperAPI
def compute_actions(self,
obs_batch,
state_batches,
state_batches=None,
prev_action_batch=None,
prev_reward_batch=None,
info_batch=None,
episodes=None,
**kwargs):
"""Compute actions for the current policy.
"""Computes actions for the current policy.
Arguments:
obs_batch (np.ndarray): batch of observations
state_batches (list): list of RNN state input batches, if any
prev_action_batch (np.ndarray): batch of previous action values
prev_reward_batch (np.ndarray): batch of previous rewards
info_batch (info): batch of info objects
Args:
obs_batch (Union[List,np.ndarray]): Batch of observations.
state_batches (Optional[list]): List of RNN state input batches,
if any.
prev_action_batch (Optional[List,np.ndarray]): Batch of previous
action values.
prev_reward_batch (Optional[List,np.ndarray]): Batch of previous
rewards.
info_batch (info): Batch of info objects.
episodes (list): MultiAgentEpisode for each obs in obs_batch.
This provides access to all of the internal episode state,
which may be useful for model-based or multiagent algorithms.
@@ -90,7 +98,7 @@ class Policy:
@DeveloperAPI
def compute_single_action(self,
obs,
state,
state=None,
prev_action=None,
prev_reward=None,
info=None,
@@ -100,10 +108,10 @@ class Policy:
"""Unbatched version of compute_actions.
Arguments:
obs (obj): single observation
state_batches (list): list of RNN state inputs, if any
prev_action (obj): previous action value, if any
prev_reward (int): previous reward, if any
obs (obj): Single observation.
state (list): List of RNN state inputs, if any.
prev_action (obj): Previous action value, if any.
prev_reward (float): Previous reward, if any.
info (dict): info object, if any
episode (MultiAgentEpisode): this provides access to all of the
internal episode state, which may be useful for model-based or
@@ -116,7 +124,6 @@ class Policy:
state_outs (list): list of RNN state outputs, if any
info (dict): dictionary of extra features, if any
"""
prev_action_batch = None
prev_reward_batch = None
info_batch = None
@@ -129,6 +136,7 @@ class Policy:
info_batch = [info]
if episode is not None:
episodes = [episode]
[action], state_out, info = self.compute_actions(
[obs], [[s] for s in state],
prev_action_batch=prev_action_batch,
@@ -137,6 +145,8 @@ class Policy:
episodes=episodes)
if clip_actions:
action = clip_action(action, self.action_space)
# Return action, internal state(s), infos.
return action, [s[0] for s in state_out], \
{k: v[0] for k, v in info.items()}
@@ -161,7 +171,7 @@ class Policy:
multi-agent algorithms.
Returns:
SampleBatch: postprocessed sample batch.
SampleBatch: Postprocessed sample batch.
"""
return sample_batch
@@ -211,7 +221,7 @@ class Policy:
Returns:
weights (obj): Serializable copy or view of model weights
"""
raise NotImplementedError
pass
@DeveloperAPI
def set_weights(self, weights):
@@ -220,7 +230,15 @@ class Policy:
Arguments:
weights (obj): Serializable copy or view of model weights
"""
raise NotImplementedError
pass
@DeveloperAPI
def num_state_tensors(self):
"""
Returns:
int: The number of RNN hidden states kept by this Policy's Model.
"""
return 0
@DeveloperAPI
def get_initial_state(self):
@@ -274,7 +292,8 @@ class Policy:
def clip_action(action, space):
"""Called to clip actions to the specified range of this policy.
"""
Called to clip actions to the specified range of this policy.
Arguments:
action: Single action.
@@ -288,8 +307,9 @@ def clip_action(action, space):
return np.clip(action, space.low, space.high)
elif isinstance(space, gym.spaces.Tuple):
if type(action) not in (tuple, list):
raise ValueError("Expected tuple space for actions {}: {}".format(
action, space))
raise ValueError(
"Expected tuple space for actions {}: {}".
format(action, space))
out = []
for a, s in zip(action, space.spaces):
out.append(clip_action(a, s))
rllib/policy/tests/__init__.py
View File
rllib/policy/tests/test_policy.py
+21
View File
@@ -0,0 +1,21 @@
import random
from ray.rllib.policy.policy import Policy
class TestPolicy(Policy):
"""
A dummy Policy that returns a random (batched) int for compute_actions
and implements all other abstract methods of Policy with "pass".
"""
def compute_actions(self,
obs_batch,
state_batches=None,
prev_action_batch=None,
prev_reward_batch=None,
episodes=None,
deterministic=None,
explore=True,
time_step=None,
**kwargs):
return [random.choice([0, 1])] * len(obs_batch), [], {}
rllib/policy/tf_policy.py
+11 -7
View File
@@ -5,7 +5,8 @@ import os
import numpy as np
import ray
import ray.experimental.tf_utils
from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY
from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY, \
ACTION_PROB, ACTION_LOGP
from ray.rllib.policy.rnn_sequencing import chop_into_sequences
from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.models.modelv2 import ModelV2
@@ -18,9 +19,6 @@ from ray.rllib.utils import try_import_tf
tf = try_import_tf()
logger = logging.getLogger(__name__)
ACTION_PROB = "action_prob"
ACTION_LOGP = "action_logp"
@DeveloperAPI
class TFPolicy(Policy):
@@ -52,6 +50,7 @@ class TFPolicy(Policy):
def __init__(self,
observation_space,
action_space,
config,
sess,
obs_input,
action_sampler,
@@ -102,9 +101,7 @@ class TFPolicy(Policy):
applying gradients. Otherwise we run all update ops found in
the current variable scope.
"""
self.observation_space = observation_space
self.action_space = action_space
super(TFPolicy, self).__init__(observation_space, action_space, config)
self.model = model
self._sess = sess
self._obs_input = obs_input
@@ -296,6 +293,13 @@ class TFPolicy(Policy):
Optional, only required to work with the multi-GPU optimizer."""
raise NotImplementedError
def is_recurrent(self):
return len(self._state_inputs) > 0
@override(Policy)
def num_state_tensors(self):
return len(self._state_inputs)
@DeveloperAPI
def extra_compute_action_feed_dict(self):
"""Extra dict to pass to the compute actions session run."""
rllib/policy/tf_policy_template.py
+2 -5
View File
@@ -96,7 +96,6 @@ def build_tf_policy(name,
Returns:
a DynamicTFPolicy instance that uses the specified args
"""
original_kwargs = locals().copy()
base = add_mixins(DynamicTFPolicy, mixins)
@@ -188,16 +187,14 @@ def build_tf_policy(name,
else:
return base.extra_compute_grad_fetches(self)
@staticmethod
def with_updates(**overrides):
return build_tf_policy(**dict(original_kwargs, **overrides))
@staticmethod
def as_eager():
return eager_tf_policy.build_eager_tf_policy(**original_kwargs)
policy_cls.with_updates = with_updates
policy_cls.as_eager = as_eager
policy_cls.with_updates = staticmethod(with_updates)
policy_cls.as_eager = staticmethod(as_eager)
policy_cls.__name__ = name
policy_cls.__qualname__ = name
return policy_cls
rllib/policy/torch_policy.py
+92 -15
View File
@@ -1,13 +1,13 @@
import numpy as np
try:
import torch
except ImportError:
pass # soft dep
from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY
from ray.rllib.utils.annotations import override
from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.utils import try_import_torch
from ray.rllib.utils.annotations import override, DeveloperAPI
from ray.rllib.utils.tracking_dict import UsageTrackingDict
from ray.rllib.utils.schedules import ConstantSchedule, PiecewiseSchedule
torch, _ = try_import_torch()
class TorchPolicy(Policy):
@@ -22,8 +22,7 @@ class TorchPolicy(Policy):
model (TorchModel): Torch model instance
dist_class (type): Torch action distribution class
"""
def __init__(self, observation_space, action_space, model, loss,
def __init__(self, observation_space, action_space, config, model, loss,
action_distribution_class):
"""Build a policy from policy and loss torch modules.
@@ -33,6 +32,7 @@ class TorchPolicy(Policy):
Arguments:
observation_space (gym.Space): observation space of the policy.
action_space (gym.Space): action space of the policy.
config (dict): The Policy config dict.
model (nn.Module): PyTorch policy module. Given observations as
input, this module must return a list of outputs where the
first item is action logits, and the rest can be any value.
@@ -41,8 +41,9 @@ class TorchPolicy(Policy):
action_distribution_class (ActionDistribution): Class for action
distribution.
"""
self.observation_space = observation_space
self.action_space = action_space
super(TorchPolicy, self).__init__(
observation_space, action_space, config
)
self.device = (torch.device("cuda")
if torch.cuda.is_available() else torch.device("cpu"))
self.model = model.to(self.device)
@@ -61,12 +62,12 @@ class TorchPolicy(Policy):
**kwargs):
with torch.no_grad():
input_dict = self._lazy_tensor_dict({
"obs": obs_batch,
SampleBatch.CUR_OBS: obs_batch,
})
if prev_action_batch:
input_dict["prev_actions"] = prev_action_batch
input_dict[SampleBatch.PREV_ACTIONS] = prev_action_batch
if prev_reward_batch:
input_dict["prev_rewards"] = prev_reward_batch
input_dict[SampleBatch.PREV_REWARDS] = prev_reward_batch
model_out = self.model(input_dict, state_batches, [1])
logits, state = model_out
action_dist = self.dist_class(logits, self.model)
@@ -128,6 +129,10 @@ class TorchPolicy(Policy):
def set_weights(self, weights):
self.model.load_state_dict(weights)
@override(Policy)
def num_state_tensors(self):
return len(self.model.get_initial_state())
@override(Policy)
def get_initial_state(self):
return [s.numpy() for s in self.model.get_initial_state()]
@@ -137,13 +142,17 @@ class TorchPolicy(Policy):
return processing info."""
return {}
def extra_action_out(self, input_dict, state_batches, model):
def extra_action_out(self, input_dict, state_batches, model,
action_dist=None):
"""Returns dict of extra info to include in experience batch.
Arguments:
input_dict (dict): Dict of model input tensors.
state_batches (list): List of state tensors.
model (TorchModelV2): Reference to the model."""
model (TorchModelV2): Reference to the model.
action_dist (Distribution): Torch Distribution object to get
log-probs (e.g. for already sampled actions).
"""
return {}
def extra_grad_info(self, train_batch):
@@ -170,3 +179,71 @@ class TorchPolicy(Policy):
train_batch.set_get_interceptor(convert)
return train_batch
@override(Policy)
def export_model(self, export_dir):
"""TODO: implement for torch.
"""
raise NotImplementedError
@override(Policy)
def export_checkpoint(self, export_dir):
"""TODO: implement for torch.
"""
raise NotImplementedError
@DeveloperAPI
class LearningRateSchedule(object):
"""Mixin for TFPolicy that adds a learning rate schedule."""
@DeveloperAPI
def __init__(self, lr, lr_schedule):
self.cur_lr = lr
if lr_schedule is None:
self.lr_schedule = ConstantSchedule(lr)
else:
self.lr_schedule = PiecewiseSchedule(
lr_schedule, outside_value=lr_schedule[-1][-1]
)
@override(Policy)
def on_global_var_update(self, global_vars):
super(LearningRateSchedule, self).on_global_var_update(global_vars)
self.cur_lr = self.lr_schedule.value(global_vars["timestep"])
@override(TorchPolicy)
def optimizer(self):
for p in self._optimizer.param_groups:
p["lr"] = self.cur_lr
return self._optimizer
@DeveloperAPI
class EntropyCoeffSchedule(object):
"""Mixin for TorchPolicy that adds entropy coeff decay."""
@DeveloperAPI
def __init__(self, entropy_coeff, entropy_coeff_schedule):
self.entropy_coeff = entropy_coeff
if entropy_coeff_schedule is None:
self.entropy_coeff_schedule = ConstantSchedule(entropy_coeff)
else:
# Allows for custom schedule similar to lr_schedule format
if isinstance(entropy_coeff_schedule, list):
self.entropy_coeff_schedule = PiecewiseSchedule(
entropy_coeff_schedule,
outside_value=entropy_coeff_schedule[-1][-1])
else:
# Implements previous version but enforces outside_value
self.entropy_coeff_schedule = PiecewiseSchedule(
[[0, entropy_coeff], [entropy_coeff_schedule, 0.0]],
outside_value=0.0)
@override(Policy)
def on_global_var_update(self, global_vars):
super(EntropyCoeffSchedule, self).on_global_var_update(global_vars)
self.entropy_coeff = self.entropy_coeff_schedule.value(
global_vars["timestep"]
)
rllib/policy/torch_policy_template.py
+13 -9
View File
@@ -77,8 +77,10 @@ def build_torch_policy(name,
self.config["model"],
framework="torch")
TorchPolicy.__init__(self, obs_space, action_space, self.model,
loss_fn, self.dist_class)
TorchPolicy.__init__(
self, obs_space, action_space, config, self.model,
loss_fn, self.dist_class
)
if after_init:
after_init(self, obs_space, action_space, config)
@@ -101,13 +103,16 @@ def build_torch_policy(name,
return TorchPolicy.extra_grad_process(self)
@override(TorchPolicy)
def extra_action_out(self, input_dict, state_batches, model):
def extra_action_out(self, input_dict, state_batches, model,
action_dist=None):
if extra_action_out_fn:
return extra_action_out_fn(self, input_dict, state_batches,
model)
return extra_action_out_fn(
self, input_dict, state_batches, model, action_dist
)
else:
return TorchPolicy.extra_action_out(self, input_dict,
state_batches, model)
return TorchPolicy.extra_action_out(
self, input_dict, state_batches, model, action_dist
)
@override(TorchPolicy)
def optimizer(self):
@@ -123,11 +128,10 @@ def build_torch_policy(name,
else:
return TorchPolicy.extra_grad_info(self, train_batch)
@staticmethod
def with_updates(**overrides):
return build_torch_policy(**dict(original_kwargs, **overrides))
policy_cls.with_updates = with_updates
policy_cls.with_updates = staticmethod(with_updates)
policy_cls.__name__ = name
policy_cls.__qualname__ = name
return policy_cls
rllib/tests/test_multi_agent_env.py
+9 -9
View File
@@ -438,13 +438,13 @@ class TestMultiAgentEnv(unittest.TestCase):
self.assertEqual(batch.policy_batches["p0"]["t"].tolist()[:10],
[4, 9, 14, 19, 24, 5, 10, 15, 20, 25])
def testCustomRNNStateValues(self):
def test_custom_rnn_state_values(self):
h = {"some": {"arbitrary": "structure", "here": [1, 2, 3]}}
class StatefulPolicy(Policy):
def compute_actions(self,
obs_batch,
state_batches,
state_batches=None,
prev_action_batch=None,
prev_reward_batch=None,
episodes=None,
@@ -465,7 +465,7 @@ class TestMultiAgentEnv(unittest.TestCase):
self.assertEqual(batch["state_in_0"][1], h)
self.assertEqual(batch["state_out_0"][1], h)
def testReturningModelBasedRolloutsData(self):
def test_returning_model_based_rollouts_data(self):
class ModelBasedPolicy(PGTFPolicy):
def compute_actions(self,
obs_batch,
@@ -512,7 +512,7 @@ class TestMultiAgentEnv(unittest.TestCase):
self.assertEqual(batch.policy_batches["p0"].count, 10)
self.assertEqual(batch.policy_batches["p1"].count, 25)
def testTrainMultiCartpoleSinglePolicy(self):
def test_train_multi_cartpole_single_policy(self):
n = 10
register_env("multi_cartpole", lambda _: MultiCartpole(n))
pg = PGTrainer(env="multi_cartpole", config={"num_workers": 0})
@@ -524,7 +524,7 @@ class TestMultiAgentEnv(unittest.TestCase):
return
raise Exception("failed to improve reward")
def testTrainMultiCartpoleMultiPolicy(self):
def test_train_multi_cartpole_multi_policy(self):
n = 10
register_env("multi_cartpole", lambda _: MultiCartpole(n))
single_env = gym.make("CartPole-v0")
@@ -625,16 +625,16 @@ class TestMultiAgentEnv(unittest.TestCase):
print(result)
raise Exception("failed to improve reward")
def testMultiAgentSyncOptimizer(self):
def test_multi_agent_sync_optimizer(self):
self._testWithOptimizer(SyncSamplesOptimizer)
def testMultiAgentAsyncGradientsOptimizer(self):
def test_multi_agent_async_gradients_optimizer(self):
self._testWithOptimizer(AsyncGradientsOptimizer)
def testMultiAgentReplayOptimizer(self):
def test_multi_agent_replay_optimizer(self):
self._testWithOptimizer(SyncReplayOptimizer)
def testTrainMultiCartpoleManyPolicies(self):
def test_train_multi_cartpole_many_policies(self):
n = 20
env = gym.make("CartPole-v0")
act_space = env.action_space
rllib/tests/test_rollout_worker.py
+23 -30
View File
@@ -20,7 +20,7 @@ from ray.tune.registry import register_env
class MockPolicy(Policy):
def compute_actions(self,
obs_batch,
state_batches,
state_batches=None,
prev_action_batch=None,
prev_reward_batch=None,
episodes=None,
@@ -35,23 +35,16 @@ class MockPolicy(Policy):
return compute_advantages(batch, 100.0, 0.9, use_gae=False)
class BadPolicy(Policy):
class BadPolicy(MockPolicy):
def compute_actions(self,
obs_batch,
state_batches,
state_batches=None,
prev_action_batch=None,
prev_reward_batch=None,
episodes=None,
**kwargs):
raise Exception("intentional error")
def postprocess_trajectory(self,
batch,
other_agent_batches=None,
episode=None):
assert episode is not None
return compute_advantages(batch, 100.0, 0.9, use_gae=False)
class FailOnStepEnv(gym.Env):
def __init__(self):
@@ -126,7 +119,7 @@ class MockVectorEnv(VectorEnv):
class TestRolloutWorker(unittest.TestCase):
def testBasic(self):
def test_basic(self):
ev = RolloutWorker(
env_creator=lambda _: gym.make("CartPole-v0"), policy=MockPolicy)
batch = ev.sample()
@@ -150,7 +143,7 @@ class TestRolloutWorker(unittest.TestCase):
to_prev(batch["actions"]))
self.assertGreater(batch["advantages"][0], 1)
def testBatchIds(self):
def test_batch_ids(self):
ev = RolloutWorker(
env_creator=lambda _: gym.make("CartPole-v0"), policy=MockPolicy)
batch1 = ev.sample()
@@ -160,7 +153,7 @@ class TestRolloutWorker(unittest.TestCase):
self.assertEqual(
len(set(SampleBatch.concat(batch1, batch2)["unroll_id"])), 2)
def testGlobalVarsUpdate(self):
def test_global_vars_update(self):
agent = A2CTrainer(
env="CartPole-v0",
config={
@@ -171,12 +164,12 @@ class TestRolloutWorker(unittest.TestCase):
result2 = agent.train()
self.assertLess(result2["info"]["learner"]["cur_lr"], 0.0001)
def testNoStepOnInit(self):
def test_no_step_on_init(self):
register_env("fail", lambda _: FailOnStepEnv())
pg = PGTrainer(env="fail", config={"num_workers": 1})
self.assertRaises(Exception, lambda: pg.train())
def testCallbacks(self):
def test_callbacks(self):
counts = Counter()
pg = PGTrainer(
env="CartPole-v0", config={
@@ -200,7 +193,7 @@ class TestRolloutWorker(unittest.TestCase):
self.assertGreater(counts["step"], 200)
self.assertLess(counts["step"], 400)
def testQueryEvaluators(self):
def test_query_evaluators(self):
register_env("test", lambda _: gym.make("CartPole-v0"))
pg = PGTrainer(
env="test",
@@ -218,7 +211,7 @@ class TestRolloutWorker(unittest.TestCase):
self.assertEqual(results2, [(0, 10), (1, 10), (2, 10)])
self.assertEqual(results3, [[1, 1], [1, 1], [1, 1]])
def testRewardClipping(self):
def test_reward_clipping(self):
# clipping on
ev = RolloutWorker(
env_creator=lambda _: MockEnv2(episode_length=10),
@@ -239,7 +232,7 @@ class TestRolloutWorker(unittest.TestCase):
result2 = collect_metrics(ev2, [])
self.assertEqual(result2["episode_reward_mean"], 1000)
def testHardHorizon(self):
def test_hard_horizon(self):
ev = RolloutWorker(
env_creator=lambda _: MockEnv(episode_length=10),
policy=MockPolicy,
@@ -253,7 +246,7 @@ class TestRolloutWorker(unittest.TestCase):
# 3 done values
self.assertEqual(sum(samples["dones"]), 3)
def testSoftHorizon(self):
def test_soft_horizon(self):
ev = RolloutWorker(
env_creator=lambda _: MockEnv(episode_length=10),
policy=MockPolicy,
@@ -267,7 +260,7 @@ class TestRolloutWorker(unittest.TestCase):
# only 1 hard done value
self.assertEqual(sum(samples["dones"]), 1)
def testMetrics(self):
def test_metrics(self):
ev = RolloutWorker(
env_creator=lambda _: MockEnv(episode_length=10),
policy=MockPolicy,
@@ -282,7 +275,7 @@ class TestRolloutWorker(unittest.TestCase):
self.assertEqual(result["episodes_this_iter"], 20)
self.assertEqual(result["episode_reward_mean"], 10)
def testAsync(self):
def test_async(self):
ev = RolloutWorker(
env_creator=lambda _: gym.make("CartPole-v0"),
sample_async=True,
@@ -292,7 +285,7 @@ class TestRolloutWorker(unittest.TestCase):
self.assertIn(key, batch)
self.assertGreater(batch["advantages"][0], 1)
def testAutoVectorization(self):
def test_auto_vectorization(self):
ev = RolloutWorker(
env_creator=lambda cfg: MockEnv(episode_length=20, config=cfg),
policy=MockPolicy,
@@ -315,7 +308,7 @@ class TestRolloutWorker(unittest.TestCase):
indices.append(env.unwrapped.config.vector_index)
self.assertEqual(indices, [0, 1, 2, 3, 4, 5, 6, 7])
def testBatchesLargerWhenVectorized(self):
def test_batches_larger_when_vectorized(self):
ev = RolloutWorker(
env_creator=lambda _: MockEnv(episode_length=8),
policy=MockPolicy,
@@ -330,7 +323,7 @@ class TestRolloutWorker(unittest.TestCase):
result = collect_metrics(ev, [])
self.assertEqual(result["episodes_this_iter"], 4)
def testVectorEnvSupport(self):
def test_vector_env_support(self):
ev = RolloutWorker(
env_creator=lambda _: MockVectorEnv(episode_length=20, num_envs=8),
policy=MockPolicy,
@@ -347,7 +340,7 @@ class TestRolloutWorker(unittest.TestCase):
result = collect_metrics(ev, [])
self.assertEqual(result["episodes_this_iter"], 8)
def testTruncateEpisodes(self):
def test_truncate_episodes(self):
ev = RolloutWorker(
env_creator=lambda _: MockEnv(10),
policy=MockPolicy,
@@ -356,7 +349,7 @@ class TestRolloutWorker(unittest.TestCase):
batch = ev.sample()
self.assertEqual(batch.count, 15)
def testCompleteEpisodes(self):
def test_complete_episodes(self):
ev = RolloutWorker(
env_creator=lambda _: MockEnv(10),
policy=MockPolicy,
@@ -365,7 +358,7 @@ class TestRolloutWorker(unittest.TestCase):
batch = ev.sample()
self.assertEqual(batch.count, 10)
def testCompleteEpisodesPacking(self):
def test_complete_episodes_packing(self):
ev = RolloutWorker(
env_creator=lambda _: MockEnv(10),
policy=MockPolicy,
@@ -377,7 +370,7 @@ class TestRolloutWorker(unittest.TestCase):
batch["t"].tolist(),
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
def testFilterSync(self):
def test_filter_sync(self):
ev = RolloutWorker(
env_creator=lambda _: gym.make("CartPole-v0"),
policy=MockPolicy,
@@ -390,7 +383,7 @@ class TestRolloutWorker(unittest.TestCase):
self.assertNotEqual(obs_f.rs.n, 0)
self.assertNotEqual(obs_f.buffer.n, 0)
def testGetFilters(self):
def test_get_filters(self):
ev = RolloutWorker(
env_creator=lambda _: gym.make("CartPole-v0"),
policy=MockPolicy,
@@ -405,7 +398,7 @@ class TestRolloutWorker(unittest.TestCase):
self.assertGreaterEqual(obs_f2.rs.n, obs_f.rs.n)
self.assertGreaterEqual(obs_f2.buffer.n, obs_f.buffer.n)
def testSyncFilter(self):
def test_sync_filter(self):
ev = RolloutWorker(
env_creator=lambda _: gym.make("CartPole-v0"),
policy=MockPolicy,