diff --git a/rllib/agents/a3c/a3c_torch_policy.py b/rllib/agents/a3c/a3c_torch_policy.py index 969781492..d4acaea74 100644 --- a/rllib/agents/a3c/a3c_torch_policy.py +++ b/rllib/agents/a3c/a3c_torch_policy.py @@ -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()} diff --git a/rllib/agents/ddpg/ddpg_policy.py b/rllib/agents/ddpg/ddpg_policy.py index 707e1d57e..6d23e8099 100644 --- a/rllib/agents/ddpg/ddpg_policy.py +++ b/rllib/agents/ddpg/ddpg_policy.py @@ -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, diff --git a/rllib/agents/dqn/apex.py b/rllib/agents/dqn/apex.py index 4aae2657c..61774436a 100644 --- a/rllib/agents/dqn/apex.py +++ b/rllib/agents/dqn/apex.py @@ -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) diff --git a/rllib/agents/marwil/marwil_policy.py b/rllib/agents/marwil/marwil_policy.py index eaae6d47b..f9a0c80d1 100644 --- a/rllib/agents/marwil/marwil_policy.py +++ b/rllib/agents/marwil/marwil_policy.py @@ -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, diff --git a/rllib/agents/ppo/ppo.py b/rllib/agents/ppo/ppo.py index 8ca45ac6d..9e96a7f8b 100644 --- a/rllib/agents/ppo/ppo.py +++ b/rllib/agents/ppo/ppo.py @@ -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): diff --git a/rllib/agents/qmix/qmix_policy.py b/rllib/agents/qmix/qmix_policy.py index f1db486d3..2464ca4d7 100644 --- a/rllib/agents/qmix/qmix_policy.py +++ b/rllib/agents/qmix/qmix_policy.py @@ -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"] diff --git a/rllib/contrib/alpha_zero/core/alpha_zero_policy.py b/rllib/contrib/alpha_zero/core/alpha_zero_policy.py index 4ed7adbb6..e002f95c8 100644 --- a/rllib/contrib/alpha_zero/core/alpha_zero_policy.py +++ b/rllib/contrib/alpha_zero/core/alpha_zero_policy.py @@ -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 diff --git a/rllib/contrib/alpha_zero/core/alpha_zero_trainer.py b/rllib/contrib/alpha_zero/core/alpha_zero_trainer.py index d523195d9..a82ee0644 100644 --- a/rllib/contrib/alpha_zero/core/alpha_zero_trainer.py +++ b/rllib/contrib/alpha_zero/core/alpha_zero_trainer.py @@ -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( diff --git a/rllib/contrib/maddpg/maddpg_policy.py b/rllib/contrib/maddpg/maddpg_policy.py index 6258b6038..50aa88246 100644 --- a/rllib/contrib/maddpg/maddpg_policy.py +++ b/rllib/contrib/maddpg/maddpg_policy.py @@ -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( diff --git a/rllib/examples/multiagent_custom_policy.py b/rllib/examples/multiagent_custom_policy.py index 54b1506f1..3342ff718 100644 --- a/rllib/examples/multiagent_custom_policy.py +++ b/rllib/examples/multiagent_custom_policy.py @@ -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, diff --git a/rllib/examples/rock_paper_scissors_multiagent.py b/rllib/examples/rock_paper_scissors_multiagent.py index fbfa92aab..614eccd2d 100644 --- a/rllib/examples/rock_paper_scissors_multiagent.py +++ b/rllib/examples/rock_paper_scissors_multiagent.py @@ -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, diff --git a/rllib/examples/rollout_worker_custom_workflow.py b/rllib/examples/rollout_worker_custom_workflow.py index 837a7feb7..316495c0b 100644 --- a/rllib/examples/rollout_worker_custom_workflow.py +++ b/rllib/examples/rollout_worker_custom_workflow.py @@ -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, diff --git a/rllib/policy/dynamic_tf_policy.py b/rllib/policy/dynamic_tf_policy.py index 0398145e3..c18c4c617 100644 --- a/rllib/policy/dynamic_tf_policy.py +++ b/rllib/policy/dynamic_tf_policy.py @@ -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() diff --git a/rllib/policy/eager_tf_policy.py b/rllib/policy/eager_tf_policy.py index f1ecb5b85..dfbe8c1bb 100644 --- a/rllib/policy/eager_tf_policy.py +++ b/rllib/policy/eager_tf_policy.py @@ -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) diff --git a/rllib/policy/policy.py b/rllib/policy/policy.py index 954a0366c..37e1ac2b5 100644 --- a/rllib/policy/policy.py +++ b/rllib/policy/policy.py @@ -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)) diff --git a/rllib/policy/tests/__init__.py b/rllib/policy/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/rllib/policy/tests/test_policy.py b/rllib/policy/tests/test_policy.py new file mode 100644 index 000000000..a34d75f18 --- /dev/null +++ b/rllib/policy/tests/test_policy.py @@ -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), [], {} diff --git a/rllib/policy/tf_policy.py b/rllib/policy/tf_policy.py index 7568fe433..0d9522129 100644 --- a/rllib/policy/tf_policy.py +++ b/rllib/policy/tf_policy.py @@ -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.""" diff --git a/rllib/policy/tf_policy_template.py b/rllib/policy/tf_policy_template.py index c7730f09e..d5387ba20 100644 --- a/rllib/policy/tf_policy_template.py +++ b/rllib/policy/tf_policy_template.py @@ -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 diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py index 5bde6b915..ab9efa692 100644 --- a/rllib/policy/torch_policy.py +++ b/rllib/policy/torch_policy.py @@ -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"] + ) diff --git a/rllib/policy/torch_policy_template.py b/rllib/policy/torch_policy_template.py index ef49f53e2..2ed1495a9 100644 --- a/rllib/policy/torch_policy_template.py +++ b/rllib/policy/torch_policy_template.py @@ -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 diff --git a/rllib/tests/test_multi_agent_env.py b/rllib/tests/test_multi_agent_env.py index 835560417..89d90de9d 100644 --- a/rllib/tests/test_multi_agent_env.py +++ b/rllib/tests/test_multi_agent_env.py @@ -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 diff --git a/rllib/tests/test_rollout_worker.py b/rllib/tests/test_rollout_worker.py index 345c33a6a..fc5914f5a 100644 --- a/rllib/tests/test_rollout_worker.py +++ b/rllib/tests/test_rollout_worker.py @@ -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,