diff --git a/rllib/agents/dqn/__init__.py b/rllib/agents/dqn/__init__.py index aaf169e3c..9572b754c 100644 --- a/rllib/agents/dqn/__init__.py +++ b/rllib/agents/dqn/__init__.py @@ -1,5 +1,7 @@ from ray.rllib.agents.dqn.apex import ApexTrainer from ray.rllib.agents.dqn.dqn import DQNTrainer, DEFAULT_CONFIG +from ray.rllib.agents.dqn.dqn_tf_policy import DQNTFPolicy +from ray.rllib.agents.dqn.dqn_torch_policy import DQNTorchPolicy from ray.rllib.agents.dqn.simple_q import SimpleQTrainer, \ DEFAULT_CONFIG as SIMPLE_Q_DEFAULT_CONFIG from ray.rllib.agents.dqn.simple_q_tf_policy import SimpleQTFPolicy @@ -7,6 +9,8 @@ from ray.rllib.agents.dqn.simple_q_torch_policy import SimpleQTorchPolicy __all__ = [ "ApexTrainer", + "DQNTFPolicy", + "DQNTorchPolicy", "DQNTrainer", "DEFAULT_CONFIG", "SIMPLE_Q_DEFAULT_CONFIG", diff --git a/rllib/agents/dqn/dqn_torch_model.py b/rllib/agents/dqn/dqn_torch_model.py index 5fae05a84..f753b5fc7 100644 --- a/rllib/agents/dqn/dqn_torch_model.py +++ b/rllib/agents/dqn/dqn_torch_model.py @@ -87,9 +87,9 @@ class DQNTorchModel(TorchModelV2): advantage_module.add_module("A", nn.Linear(ins, action_space.n)) value_module.add_module("V", nn.Linear(ins, 1)) # Non-dueling: - # Q-value layer (use Advantage module's outputs as Q-values). + # Q-value layer (use main module's outputs as Q-values). else: - advantage_module.add_module("Q", nn.Linear(ins, action_space.n)) + pass self.advantage_module = advantage_module self.value_module = value_module diff --git a/rllib/agents/dqn/dqn_torch_policy.py b/rllib/agents/dqn/dqn_torch_policy.py index 7050b3cfa..22bf7905f 100644 --- a/rllib/agents/dqn/dqn_torch_policy.py +++ b/rllib/agents/dqn/dqn_torch_policy.py @@ -19,7 +19,7 @@ from ray.rllib.utils import try_import_torch torch, nn = try_import_torch() F = None if nn: - F = torch.nn.functional + F = nn.functional class QLoss: diff --git a/rllib/agents/pg/__init__.py b/rllib/agents/pg/__init__.py index b8d289645..b2592a6c6 100644 --- a/rllib/agents/pg/__init__.py +++ b/rllib/agents/pg/__init__.py @@ -1,9 +1,14 @@ from ray.rllib.agents.pg.pg import PGTrainer, DEFAULT_CONFIG from ray.rllib.agents.pg.pg_tf_policy import pg_tf_loss, \ - post_process_advantages -from ray.rllib.agents.pg.pg_torch_policy import pg_torch_loss + post_process_advantages, PGTFPolicy +from ray.rllib.agents.pg.pg_torch_policy import pg_torch_loss, PGTorchPolicy __all__ = [ - "PGTrainer", "pg_tf_loss", "pg_torch_loss", "post_process_advantages", - "DEFAULT_CONFIG" + "pg_tf_loss", + "pg_torch_loss", + "post_process_advantages", + "DEFAULT_CONFIG", + "PGTFPolicy", + "PGTorchPolicy", + "PGTrainer", ] diff --git a/rllib/agents/ppo/__init__.py b/rllib/agents/ppo/__init__.py index 741ed15c7..120773707 100644 --- a/rllib/agents/ppo/__init__.py +++ b/rllib/agents/ppo/__init__.py @@ -1,5 +1,14 @@ from ray.rllib.agents.ppo.ppo import PPOTrainer, DEFAULT_CONFIG +from ray.rllib.agents.ppo.ppo_tf_policy import PPOTFPolicy +from ray.rllib.agents.ppo.ppo_torch_policy import PPOTorchPolicy from ray.rllib.agents.ppo.appo import APPOTrainer from ray.rllib.agents.ppo.ddppo import DDPPOTrainer -__all__ = ["APPOTrainer", "DDPPOTrainer", "PPOTrainer", "DEFAULT_CONFIG"] +__all__ = [ + "APPOTrainer", + "DDPPOTrainer", + "DEFAULT_CONFIG", + "PPOTFPolicy", + "PPOTorchPolicy", + "PPOTrainer", +] diff --git a/rllib/agents/qmix/model.py b/rllib/agents/qmix/model.py index 66ec26975..654f62648 100644 --- a/rllib/agents/qmix/model.py +++ b/rllib/agents/qmix/model.py @@ -23,7 +23,7 @@ class RNNModel(TorchModelV2, nn.Module): @override(ModelV2) def get_initial_state(self): - # make hidden states on same device as model + # Place hidden states on same device as model. return [self.fc1.weight.new(1, self.rnn_hidden_dim).zero_().squeeze(0)] @override(ModelV2) diff --git a/rllib/examples/env/correlated_actions_env.py b/rllib/examples/env/correlated_actions_env.py index 4a42c5ca1..db0c78e01 100644 --- a/rllib/examples/env/correlated_actions_env.py +++ b/rllib/examples/env/correlated_actions_env.py @@ -11,21 +11,23 @@ class CorrelatedActionsEnv(gym.Env): def __init__(self, _): self.observation_space = Discrete(2) self.action_space = Tuple([Discrete(2), Discrete(2)]) + self.last_observation = None def reset(self): self.t = 0 - self.last = random.choice([0, 1]) - return self.last + self.last_observation = random.choice([0, 1]) + return self.last_observation def step(self, action): self.t += 1 a1, a2 = action reward = 0 - if a1 == self.last: + # Encourage correlation between most recent observation and a1. + if a1 == self.last_observation: reward += 5 - # encourage correlation between a1 and a2 + # Encourage correlation between a1 and a2. if a1 == a2: reward += 5 done = self.t > 20 - self.last = random.choice([0, 1]) - return self.last, reward, done, {} + self.last_observation = random.choice([0, 1]) + return self.last_observation, reward, done, {} diff --git a/rllib/examples/env/rock_paper_scissors.py b/rllib/examples/env/rock_paper_scissors.py index a4480013c..bd27da7f0 100644 --- a/rllib/examples/env/rock_paper_scissors.py +++ b/rllib/examples/env/rock_paper_scissors.py @@ -23,6 +23,9 @@ class RockPaperScissors(MultiAgentEnv): self.last_move = None self.num_moves = 0 + # For test-case inspections (compare both players' scores). + self.player1_score = self.player2_score = 0 + def reset(self): self.last_move = (0, 0) self.num_moves = 0 @@ -79,4 +82,10 @@ class RockPaperScissors(MultiAgentEnv): done = { "__all__": self.num_moves >= 10, } + + if rew["player1"] > rew["player2"]: + self.player1_score += 1 + elif rew["player2"] > rew["player1"]: + self.player2_score += 1 + return obs, rew, done, {} diff --git a/rllib/examples/models/__init__.py b/rllib/examples/models/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/rllib/examples/models/autoregressive_action_dist.py b/rllib/examples/models/autoregressive_action_dist.py new file mode 100644 index 000000000..324913a68 --- /dev/null +++ b/rllib/examples/models/autoregressive_action_dist.py @@ -0,0 +1,147 @@ +from ray.rllib.models.tf.tf_action_dist import Categorical, ActionDistribution +from ray.rllib.models.torch.torch_action_dist import TorchCategorical, \ + TorchDistributionWrapper +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf = try_import_tf() +torch, nn = try_import_torch() + + +class BinaryAutoregressiveDistribution(ActionDistribution): + """Action distribution P(a1, a2) = P(a1) * P(a2 | a1)""" + + def deterministic_sample(self): + # first, sample a1 + a1_dist = self._a1_distribution() + a1 = a1_dist.deterministic_sample() + + # sample a2 conditioned on a1 + a2_dist = self._a2_distribution(a1) + a2 = a2_dist.deterministic_sample() + self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2) + + # return the action tuple + return (a1, a2) + + def sample(self): + # first, sample a1 + a1_dist = self._a1_distribution() + a1 = a1_dist.sample() + + # sample a2 conditioned on a1 + a2_dist = self._a2_distribution(a1) + a2 = a2_dist.sample() + self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2) + + # return the action tuple + return (a1, a2) + + def logp(self, actions): + a1, a2 = actions[:, 0], actions[:, 1] + a1_vec = tf.expand_dims(tf.cast(a1, tf.float32), 1) + a1_logits, a2_logits = self.model.action_model([self.inputs, a1_vec]) + return ( + Categorical(a1_logits).logp(a1) + Categorical(a2_logits).logp(a2)) + + def sampled_action_logp(self): + return tf.exp(self._action_logp) + + def entropy(self): + a1_dist = self._a1_distribution() + a2_dist = self._a2_distribution(a1_dist.sample()) + return a1_dist.entropy() + a2_dist.entropy() + + def kl(self, other): + a1_dist = self._a1_distribution() + a1_terms = a1_dist.kl(other._a1_distribution()) + + a1 = a1_dist.sample() + a2_terms = self._a2_distribution(a1).kl(other._a2_distribution(a1)) + return a1_terms + a2_terms + + def _a1_distribution(self): + BATCH = tf.shape(self.inputs)[0] + a1_logits, _ = self.model.action_model( + [self.inputs, tf.zeros((BATCH, 1))]) + a1_dist = Categorical(a1_logits) + return a1_dist + + def _a2_distribution(self, a1): + a1_vec = tf.expand_dims(tf.cast(a1, tf.float32), 1) + _, a2_logits = self.model.action_model([self.inputs, a1_vec]) + a2_dist = Categorical(a2_logits) + return a2_dist + + @staticmethod + def required_model_output_shape(self, model_config): + return 16 # controls model output feature vector size + + +class TorchBinaryAutoregressiveDistribution(TorchDistributionWrapper): + """Action distribution P(a1, a2) = P(a1) * P(a2 | a1)""" + + def deterministic_sample(self): + # first, sample a1 + a1_dist = self._a1_distribution() + a1 = a1_dist.deterministic_sample() + + # sample a2 conditioned on a1 + a2_dist = self._a2_distribution(a1) + a2 = a2_dist.deterministic_sample() + self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2) + + # return the action tuple + return (a1, a2) + + def sample(self): + # first, sample a1 + a1_dist = self._a1_distribution() + a1 = a1_dist.sample() + + # sample a2 conditioned on a1 + a2_dist = self._a2_distribution(a1) + a2 = a2_dist.sample() + self._action_logp = a1_dist.logp(a1) + a2_dist.logp(a2) + + # return the action tuple + return (a1, a2) + + def logp(self, actions): + a1, a2 = actions[:, 0], actions[:, 1] + a1_vec = torch.unsqueeze(a1.float(), 1) + a1_logits, a2_logits = self.model.action_module(self.inputs, a1_vec) + return (TorchCategorical(a1_logits).logp(a1) + + TorchCategorical(a2_logits).logp(a2)) + + def sampled_action_logp(self): + return torch.exp(self._action_logp) + + def entropy(self): + a1_dist = self._a1_distribution() + a2_dist = self._a2_distribution(a1_dist.sample()) + return a1_dist.entropy() + a2_dist.entropy() + + def kl(self, other): + a1_dist = self._a1_distribution() + a1_terms = a1_dist.kl(other._a1_distribution()) + + a1 = a1_dist.sample() + a2_terms = self._a2_distribution(a1).kl(other._a2_distribution(a1)) + return a1_terms + a2_terms + + def _a1_distribution(self): + BATCH = self.inputs.shape[0] + a1_logits, _ = self.model.action_module(self.inputs, + torch.zeros((BATCH, 1))) + a1_dist = TorchCategorical(a1_logits) + return a1_dist + + def _a2_distribution(self, a1): + a1_vec = torch.unsqueeze(a1.float(), 1) + _, a2_logits = self.model.action_module(self.inputs, a1_vec) + a2_dist = TorchCategorical(a2_logits) + return a2_dist + + @staticmethod + def required_model_output_shape(self, model_config): + return 16 # controls model output feature vector size diff --git a/rllib/examples/models/autoregressive_action_model.py b/rllib/examples/models/autoregressive_action_model.py new file mode 100644 index 000000000..ef570a7a5 --- /dev/null +++ b/rllib/examples/models/autoregressive_action_model.py @@ -0,0 +1,157 @@ +from gym.spaces import Discrete, Tuple + +from ray.rllib.models.tf.misc import normc_initializer +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.torch.misc import normc_initializer as normc_init_torch +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf = try_import_tf() +torch, nn = try_import_torch() + + +class AutoregressiveActionModel(TFModelV2): + """Implements the `.action_model` branch required above.""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + super(AutoregressiveActionModel, self).__init__( + obs_space, action_space, num_outputs, model_config, name) + if action_space != Tuple([Discrete(2), Discrete(2)]): + raise ValueError( + "This model only supports the [2, 2] action space") + + # Inputs + obs_input = tf.keras.layers.Input( + shape=obs_space.shape, name="obs_input") + a1_input = tf.keras.layers.Input(shape=(1, ), name="a1_input") + ctx_input = tf.keras.layers.Input( + shape=(num_outputs, ), name="ctx_input") + + # Output of the model (normally 'logits', but for an autoregressive + # dist this is more like a context/feature layer encoding the obs) + context = tf.keras.layers.Dense( + num_outputs, + name="hidden", + activation=tf.nn.tanh, + kernel_initializer=normc_initializer(1.0))(obs_input) + + # V(s) + value_out = tf.keras.layers.Dense( + 1, + name="value_out", + activation=None, + kernel_initializer=normc_initializer(0.01))(context) + + # P(a1 | obs) + a1_logits = tf.keras.layers.Dense( + 2, + name="a1_logits", + activation=None, + kernel_initializer=normc_initializer(0.01))(ctx_input) + + # P(a2 | a1) + # --note: typically you'd want to implement P(a2 | a1, obs) as follows: + # a2_context = tf.keras.layers.Concatenate(axis=1)( + # [ctx_input, a1_input]) + a2_context = a1_input + a2_hidden = tf.keras.layers.Dense( + 16, + name="a2_hidden", + activation=tf.nn.tanh, + kernel_initializer=normc_initializer(1.0))(a2_context) + a2_logits = tf.keras.layers.Dense( + 2, + name="a2_logits", + activation=None, + kernel_initializer=normc_initializer(0.01))(a2_hidden) + + # Base layers + self.base_model = tf.keras.Model(obs_input, [context, value_out]) + self.register_variables(self.base_model.variables) + self.base_model.summary() + + # Autoregressive action sampler + self.action_model = tf.keras.Model([ctx_input, a1_input], + [a1_logits, a2_logits]) + self.action_model.summary() + self.register_variables(self.action_model.variables) + + def forward(self, input_dict, state, seq_lens): + context, self._value_out = self.base_model(input_dict["obs"]) + return context, state + + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class TorchAutoregressiveActionModel(TorchModelV2, nn.Module): + """PyTorch version of the AutoregressiveActionModel above.""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + TorchModelV2.__init__(self, obs_space, action_space, num_outputs, + model_config, name) + nn.Module.__init__(self) + + if action_space != Tuple([Discrete(2), Discrete(2)]): + raise ValueError( + "This model only supports the [2, 2] action space") + + # Output of the model (normally 'logits', but for an autoregressive + # dist this is more like a context/feature layer encoding the obs) + self.context_layer = SlimFC( + in_size=obs_space.shape[0], + out_size=num_outputs, + initializer=normc_init_torch(1.0), + activation_fn=nn.Tanh, + ) + + # V(s) + self.value_branch = SlimFC( + in_size=num_outputs, + out_size=1, + initializer=normc_init_torch(0.01), + activation_fn=None, + ) + + # P(a1 | obs) + self.a1_logits = SlimFC( + in_size=num_outputs, + out_size=2, + activation_fn=None, + initializer=normc_init_torch(0.01)) + + class _ActionModel(nn.Module): + def __init__(self): + nn.Module.__init__(self) + self.a2_hidden = SlimFC( + in_size=1, + out_size=16, + activation_fn=nn.Tanh, + initializer=normc_init_torch(1.0)) + self.a2_logits = SlimFC( + in_size=16, + out_size=2, + activation_fn=None, + initializer=normc_init_torch(0.01)) + + def forward(self_, ctx_input, a1_input): + a1_logits = self.a1_logits(ctx_input) + a2_logits = self_.a2_logits(self_.a2_hidden(a1_input)) + return a1_logits, a2_logits + + # P(a2 | a1) + # --note: typically you'd want to implement P(a2 | a1, obs) as follows: + # a2_context = tf.keras.layers.Concatenate(axis=1)( + # [ctx_input, a1_input]) + self.action_module = _ActionModel() + + def forward(self, input_dict, state, seq_lens): + context = self.context_layer(input_dict["obs"]) + self._value_out = self.value_branch(context) + return context, state + + def value_function(self): + return torch.reshape(self._value_out, [-1]) diff --git a/rllib/examples/models/batch_norm_model.py b/rllib/examples/models/batch_norm_model.py new file mode 100644 index 000000000..693df3910 --- /dev/null +++ b/rllib/examples/models/batch_norm_model.py @@ -0,0 +1,190 @@ +import numpy as np + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.misc import normc_initializer +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.torch.misc import SlimFC, normc_initializer as \ + torch_normc_initializer +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils import try_import_tf, try_import_torch +from ray.rllib.utils.annotations import override + +tf = try_import_tf() +torch, nn = try_import_torch() + + +class BatchNormModel(TFModelV2): + """Example of a TFModelV2 that is built w/o using tf.keras. + + NOTE: This example does not work when using a keras-based TFModelV2 due + to a bug in keras related to missing values for input placeholders, even + though these input values have been provided in a forward pass through the + actual keras Model. + + All Model logic (layers) is defined in the `forward` method (incl. + the batch_normalization layers). Also, all variables are registered + (only once) at the end of `forward`, so an optimizer knows which tensors + to train on. A standard `value_function` override is used. + """ + capture_index = 0 + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + super().__init__(obs_space, action_space, num_outputs, model_config, + name) + # Have we registered our vars yet (see `forward`)? + self._registered = False + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + last_layer = input_dict["obs"] + hiddens = [256, 256] + with tf.variable_scope("model", reuse=tf.AUTO_REUSE): + for i, size in enumerate(hiddens): + last_layer = tf.layers.dense( + last_layer, + size, + kernel_initializer=normc_initializer(1.0), + activation=tf.nn.tanh, + name="fc{}".format(i)) + # Add a batch norm layer + last_layer = tf.layers.batch_normalization( + last_layer, + training=input_dict["is_training"], + name="bn_{}".format(i)) + + output = tf.layers.dense( + last_layer, + self.num_outputs, + kernel_initializer=normc_initializer(0.01), + activation=None, + name="out") + self._value_out = tf.layers.dense( + last_layer, + 1, + kernel_initializer=normc_initializer(1.0), + activation=None, + name="vf") + if not self._registered: + self.register_variables( + tf.get_collection( + tf.GraphKeys.TRAINABLE_VARIABLES, scope=".+/model/.+")) + self._registered = True + + return output, [] + + @override(ModelV2) + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class KerasBatchNormModel(TFModelV2): + """Keras version of above BatchNormModel with exactly the same structure. + + IMORTANT NOTE: This model will not work with PPO due to a bug in keras + that surfaces when having more than one input placeholder (here: `inputs` + and `is_training`) AND using the `make_tf_callable` helper (e.g. used by + PPO), in which auto-placeholders are generated, then passed through the + tf.keras. models.Model. In this last step, the connection between 1) the + provided value in the auto-placeholder and 2) the keras `is_training` + Input is broken and keras complains. + Use the above `BatchNormModel` (a non-keras based TFModelV2), instead. + """ + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + super().__init__(obs_space, action_space, num_outputs, model_config, + name) + inputs = tf.keras.layers.Input(shape=obs_space.shape, name="inputs") + is_training = tf.keras.layers.Input( + shape=(), dtype=tf.bool, batch_size=1, name="is_training") + last_layer = inputs + hiddens = [256, 256] + for i, size in enumerate(hiddens): + label = "fc{}".format(i) + last_layer = tf.keras.layers.Dense( + units=size, + kernel_initializer=normc_initializer(1.0), + activation=tf.nn.tanh, + name=label)(last_layer) + # Add a batch norm layer + last_layer = tf.keras.layers.BatchNormalization()( + last_layer, training=is_training[0]) + output = tf.keras.layers.Dense( + units=self.num_outputs, + kernel_initializer=normc_initializer(0.01), + activation=None, + name="fc_out")(last_layer) + value_out = tf.keras.layers.Dense( + units=1, + kernel_initializer=normc_initializer(0.01), + activation=None, + name="value_out")(last_layer) + + self.base_model = tf.keras.models.Model( + inputs=[inputs, is_training], outputs=[output, value_out]) + self.register_variables(self.base_model.variables) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + out, self._value_out = self.base_model( + [input_dict["obs"], input_dict["is_training"]]) + return out, [] + + @override(ModelV2) + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class TorchBatchNormModel(TorchModelV2, nn.Module): + """Example of a TorchModelV2 using batch normalization.""" + capture_index = 0 + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name, **kwargs): + TorchModelV2.__init__(self, obs_space, action_space, num_outputs, + model_config, name) + nn.Module.__init__(self) + layers = [] + prev_layer_size = int(np.product(obs_space.shape)) + self._logits = None + + # Create layers 0 to second-last. + for size in [256, 256]: + layers.append( + SlimFC( + in_size=prev_layer_size, + out_size=size, + initializer=torch_normc_initializer(1.0), + activation_fn=nn.ReLU)) + prev_layer_size = size + # Add a batch norm layer. + layers.append(nn.BatchNorm1d(prev_layer_size)) + + self._logits = SlimFC( + in_size=prev_layer_size, + out_size=self.num_outputs, + initializer=torch_normc_initializer(0.01), + activation_fn=None) + + self._value_branch = SlimFC( + in_size=prev_layer_size, + out_size=1, + initializer=torch_normc_initializer(1.0), + activation_fn=None) + + self._hidden_layers = nn.Sequential(*layers) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + # Set the correct train-mode for our hidden module (only important + # b/c we have some batch-norm layers). + self._hidden_layers.train(mode=input_dict["is_training"]) + hidden_out = self._hidden_layers(input_dict["obs"]) + logits = self._logits(hidden_out) + self._value_out = self._value_branch(hidden_out) + return logits, [] + + @override(ModelV2) + def value_function(self): + return torch.reshape(self._value_out, [-1]) diff --git a/rllib/examples/models/centralized_critic_models.py b/rllib/examples/models/centralized_critic_models.py new file mode 100644 index 000000000..0e8937510 --- /dev/null +++ b/rllib/examples/models/centralized_critic_models.py @@ -0,0 +1,173 @@ +from gym.spaces import Box + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.tf.fcnet_v2 import FullyConnectedNetwork +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFC +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf = try_import_tf() +torch, nn = try_import_torch() + + +class CentralizedCriticModel(TFModelV2): + """Multi-agent model that implements a centralized value function.""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + super(CentralizedCriticModel, self).__init__( + obs_space, action_space, num_outputs, model_config, name) + # Base of the model + self.model = FullyConnectedNetwork(obs_space, action_space, + num_outputs, model_config, name) + self.register_variables(self.model.variables()) + + # Central VF maps (obs, opp_obs, opp_act) -> vf_pred + obs = tf.keras.layers.Input(shape=(6, ), name="obs") + opp_obs = tf.keras.layers.Input(shape=(6, ), name="opp_obs") + opp_act = tf.keras.layers.Input(shape=(2, ), name="opp_act") + concat_obs = tf.keras.layers.Concatenate(axis=1)( + [obs, opp_obs, opp_act]) + central_vf_dense = tf.keras.layers.Dense( + 16, activation=tf.nn.tanh, name="c_vf_dense")(concat_obs) + central_vf_out = tf.keras.layers.Dense( + 1, activation=None, name="c_vf_out")(central_vf_dense) + self.central_vf = tf.keras.Model( + inputs=[obs, opp_obs, opp_act], outputs=central_vf_out) + self.register_variables(self.central_vf.variables) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + return self.model.forward(input_dict, state, seq_lens) + + def central_value_function(self, obs, opponent_obs, opponent_actions): + return tf.reshape( + self.central_vf( + [obs, opponent_obs, + tf.one_hot(opponent_actions, 2)]), [-1]) + + @override(ModelV2) + def value_function(self): + return self.model.value_function() # not used + + +class YetAnotherCentralizedCriticModel(TFModelV2): + """Multi-agent model that implements a centralized value function. + + It assumes the observation is a dict with 'own_obs' and 'opponent_obs', the + former of which can be used for computing actions (i.e., decentralized + execution), and the latter for optimization (i.e., centralized learning). + + This model has two parts: + - An action model that looks at just 'own_obs' to compute actions + - A value model that also looks at the 'opponent_obs' / 'opponent_action' + to compute the value (it does this by using the 'obs_flat' tensor). + """ + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + super(YetAnotherCentralizedCriticModel, self).__init__( + obs_space, action_space, num_outputs, model_config, name) + + self.action_model = FullyConnectedNetwork( + Box(low=0, high=1, shape=(6, )), # one-hot encoded Discrete(6) + action_space, + num_outputs, + model_config, + name + "_action") + self.register_variables(self.action_model.variables()) + + self.value_model = FullyConnectedNetwork(obs_space, action_space, 1, + model_config, name + "_vf") + self.register_variables(self.value_model.variables()) + + def forward(self, input_dict, state, seq_lens): + self._value_out, _ = self.value_model({ + "obs": input_dict["obs_flat"] + }, state, seq_lens) + return self.action_model({ + "obs": input_dict["obs"]["own_obs"] + }, state, seq_lens) + + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class TorchCentralizedCriticModel(TorchModelV2, nn.Module): + """Multi-agent model that implements a centralized VF.""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + TorchModelV2.__init__(self, obs_space, action_space, num_outputs, + model_config, name) + nn.Module.__init__(self) + + # Base of the model + self.model = TorchFC(obs_space, action_space, num_outputs, + model_config, name) + + # Central VF maps (obs, opp_obs, opp_act) -> vf_pred + input_size = 6 + 6 + 2 # obs + opp_obs + opp_act + self.central_vf_dense = SlimFC(input_size, 16, activation_fn=nn.Tanh) + self.central_vf_out = SlimFC(16, 1) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + model_out, _ = self.model(input_dict, state, seq_lens) + return model_out, [] + + def central_value_function(self, obs, opponent_obs, opponent_actions): + input_ = torch.cat([ + obs, opponent_obs, + torch.nn.functional.one_hot(opponent_actions, 2) + ], 1) + return torch.reshape( + self.central_vf_out(self.central_vf_dense(input_)), [-1]) + + @override(ModelV2) + def value_function(self): + return self.model.value_function() # not used + + +class YetAnotherTorchCentralizedCriticModel(TorchModelV2, nn.Module): + """Multi-agent model that implements a centralized value function. + + It assumes the observation is a dict with 'own_obs' and 'opponent_obs', the + former of which can be used for computing actions (i.e., decentralized + execution), and the latter for optimization (i.e., centralized learning). + + This model has two parts: + - An action model that looks at just 'own_obs' to compute actions + - A value model that also looks at the 'opponent_obs' / 'opponent_action' + to compute the value (it does this by using the 'obs_flat' tensor). + """ + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + TorchModelV2.__init__(self, obs_space, action_space, num_outputs, + model_config, name) + nn.Module.__init__(self) + + self.action_model = TorchFC( + Box(low=0, high=1, shape=(6, )), # one-hot encoded Discrete(6) + action_space, + num_outputs, + model_config, + name + "_action") + + self.value_model = TorchFC(obs_space, action_space, 1, model_config, + name + "_vf") + + def forward(self, input_dict, state, seq_lens): + self._value_out, _ = self.value_model({ + "obs": input_dict["obs_flat"] + }, state, seq_lens) + return self.action_model({ + "obs": input_dict["obs"]["own_obs"] + }, state, seq_lens) + + def value_function(self): + return torch.reshape(self._value_out, [-1]) diff --git a/rllib/examples/models/eager_model.py b/rllib/examples/models/eager_model.py new file mode 100644 index 000000000..3a6ae83c6 --- /dev/null +++ b/rllib/examples/models/eager_model.py @@ -0,0 +1,60 @@ +import random + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.fcnet_v2 import FullyConnectedNetwork +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf + +tf = try_import_tf() + + +class EagerModel(TFModelV2): + """Example of using embedded eager execution in a custom model. + + This shows how to use tf.py_function() to execute a snippet of TF code + in eager mode. Here the `self.forward_eager` method just prints out + the intermediate tensor for debug purposes, but you can in general + perform any TF eager operation in tf.py_function(). + """ + + def __init__(self, observation_space, action_space, num_outputs, + model_config, name): + super().__init__(observation_space, action_space, num_outputs, + model_config, name) + + inputs = tf.keras.layers.Input(shape=observation_space.shape) + self.fcnet = FullyConnectedNetwork( + obs_space=self.obs_space, + action_space=self.action_space, + num_outputs=self.num_outputs, + model_config=self.model_config, + name="fc1") + out, value_out = self.fcnet.base_model(inputs) + + def lambda_(x): + eager_out = tf.py_function(self.forward_eager, [x], tf.float32) + with tf.control_dependencies([eager_out]): + eager_out.set_shape(x.shape) + return eager_out + + out = tf.keras.layers.Lambda(lambda_)(out) + self.base_model = tf.keras.models.Model(inputs, [out, value_out]) + self.register_variables(self.base_model.variables) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + out, self._value_out = self.base_model(input_dict["obs"], state, + seq_lens) + return out, [] + + @override(ModelV2) + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + def forward_eager(self, feature_layer): + assert tf.executing_eagerly() + if random.random() > 0.99: + print("Eagerly printing the feature layer mean value", + tf.reduce_mean(feature_layer)) + return feature_layer diff --git a/rllib/examples/models/fast_model.py b/rllib/examples/models/fast_model.py new file mode 100644 index 000000000..15c35a752 --- /dev/null +++ b/rllib/examples/models/fast_model.py @@ -0,0 +1,76 @@ +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf = try_import_tf() +torch, nn = try_import_torch() + + +class FastModel(TFModelV2): + """An example for a non-Keras ModelV2 in tf that learns a single weight. + + Defines all network architecture in `forward` (not `__init__` as it's + usually done for Keras-style TFModelV2s). + """ + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + super().__init__(obs_space, action_space, num_outputs, model_config, + name) + # Have we registered our vars yet (see `forward`)? + self._registered = False + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + with tf.variable_scope("model", reuse=tf.AUTO_REUSE): + bias = tf.get_variable( + dtype=tf.float32, + name="bias", + initializer=tf.zeros_initializer, + shape=()) + output = bias + \ + tf.zeros([tf.shape(input_dict["obs"])[0], self.num_outputs]) + self._value_out = tf.reduce_mean(output, -1) # fake value + + if not self._registered: + self.register_variables( + tf.get_collection( + tf.GraphKeys.TRAINABLE_VARIABLES, scope=".+/model/.+")) + self._registered = True + + return output, [] + + @override(ModelV2) + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class TorchFastModel(TorchModelV2, nn.Module): + """Torch version of FastModel (tf).""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + TorchModelV2.__init__(self, obs_space, action_space, num_outputs, + model_config, name) + nn.Module.__init__(self) + + self.bias = torch.tensor( + [0.0], dtype=torch.float32, requires_grad=True) + + # Only needed to give some params to the optimizer (even though, + # they are never used anywhere). + self.dummy_layer = SlimFC(1, 1) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + output = self.bias + \ + torch.zeros(size=(input_dict["obs"].shape[0], self.num_outputs)) + self._value_out = torch.mean(output, -1) # fake value + return output, [] + + @override(ModelV2) + def value_function(self): + return torch.reshape(self._value_out, [-1]) diff --git a/rllib/examples/models/mobilenet_v2_with_lstm_models.py b/rllib/examples/models/mobilenet_v2_with_lstm_models.py new file mode 100644 index 000000000..ade63d73f --- /dev/null +++ b/rllib/examples/models/mobilenet_v2_with_lstm_models.py @@ -0,0 +1,150 @@ +import numpy as np + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.recurrent_tf_modelv2 import RecurrentTFModelV2 +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.recurrent_torch_model import RecurrentTorchModel +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf = try_import_tf() +torch, nn = try_import_torch() + + +class MobileV2PlusRNNModel(RecurrentTFModelV2): + """A conv. + recurrent keras net example using a pre-trained MobileNet.""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name, cnn_shape): + + super(MobileV2PlusRNNModel, self).__init__( + obs_space, action_space, num_outputs, model_config, name) + + self.cell_size = 16 + visual_size = cnn_shape[0] * cnn_shape[1] * cnn_shape[2] + + state_in_h = tf.keras.layers.Input(shape=(self.cell_size, ), name="h") + state_in_c = tf.keras.layers.Input(shape=(self.cell_size, ), name="c") + seq_in = tf.keras.layers.Input(shape=(), name="seq_in", dtype=tf.int32) + + inputs = tf.keras.layers.Input( + shape=(None, visual_size), name="visual_inputs") + + input_visual = inputs + input_visual = tf.reshape( + input_visual, [-1, cnn_shape[0], cnn_shape[1], cnn_shape[2]]) + cnn_input = tf.keras.layers.Input(shape=cnn_shape, name="cnn_input") + + cnn_model = tf.keras.applications.mobilenet_v2.MobileNetV2( + alpha=1.0, + include_top=True, + weights=None, + input_tensor=cnn_input, + pooling=None) + vision_out = cnn_model(input_visual) + vision_out = tf.reshape( + vision_out, + [-1, tf.shape(inputs)[1], + vision_out.shape.as_list()[-1]]) + + lstm_out, state_h, state_c = tf.keras.layers.LSTM( + self.cell_size, + return_sequences=True, + return_state=True, + name="lstm")( + inputs=vision_out, + mask=tf.sequence_mask(seq_in), + initial_state=[state_in_h, state_in_c]) + + # Postprocess LSTM output with another hidden layer and compute values. + logits = tf.keras.layers.Dense( + self.num_outputs, + activation=tf.keras.activations.linear, + name="logits")(lstm_out) + values = tf.keras.layers.Dense( + 1, activation=None, name="values")(lstm_out) + + # Create the RNN model + self.rnn_model = tf.keras.Model( + inputs=[inputs, seq_in, state_in_h, state_in_c], + outputs=[logits, values, state_h, state_c]) + self.register_variables(self.rnn_model.variables) + self.rnn_model.summary() + + @override(RecurrentTFModelV2) + def forward_rnn(self, inputs, state, seq_lens): + model_out, self._value_out, h, c = self.rnn_model([inputs, seq_lens] + + state) + return model_out, [h, c] + + @override(ModelV2) + def get_initial_state(self): + return [ + np.zeros(self.cell_size, np.float32), + np.zeros(self.cell_size, np.float32), + ] + + @override(ModelV2) + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class TorchMobileV2PlusRNNModel(RecurrentTorchModel): + """A conv. + recurrent torch net example using a pre-trained MobileNet.""" + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name, cnn_shape): + + super().__init__(obs_space, action_space, num_outputs, model_config, + name) + + self.lstm_state_size = 16 + self.cnn_shape = list(cnn_shape) + self.visual_size_in = cnn_shape[0] * cnn_shape[1] * cnn_shape[2] + # MobileNetV2 has a flat output of (1000,). + self.visual_size_out = 1000 + + # Load the MobileNetV2 from torch.hub. + self.cnn_model = torch.hub.load( + "pytorch/vision:v0.6.0", "mobilenet_v2", pretrained=True) + + self.lstm = nn.LSTM( + self.visual_size_out, self.lstm_state_size, batch_first=True) + + # Postprocess LSTM output with another hidden layer and compute values. + self.logits = SlimFC(self.lstm_state_size, self.num_outputs) + self.value_branch = SlimFC(self.lstm_state_size, 1) + + @override(RecurrentTFModelV2) + def forward_rnn(self, inputs, state, seq_lens): + # Create image dims. + vision_in = torch.reshape(inputs, [-1] + self.cnn_shape) + vision_out = self.cnn_model(vision_in) + # Flatten. + vision_out_time_ranked = torch.reshape( + vision_out, + [inputs.shape[0], inputs.shape[1], vision_out.shape[-1]]) + if len(state[0].shape) == 2: + state[0] = state[0].unsqueeze(0) + state[1] = state[1].unsqueeze(0) + # Forward through LSTM. + lstm_out, [h, c] = self.lstm(vision_out_time_ranked, state) + # Forward LSTM out through logits layer and value layer. + logits = self.logits(lstm_out) + self._value_out = self.value_branch(lstm_out) + return logits, [h.squeeze(0), c.squeeze(0)] + + @override(ModelV2) + def get_initial_state(self): + # Place hidden states on same device as model. + h = [ + list(self.cnn_model.modules())[-1].weight.new( + 1, self.lstm_state_size).zero_().squeeze(0), + list(self.cnn_model.modules())[-1].weight.new( + 1, self.lstm_state_size).zero_().squeeze(0), + ] + return h + + @override(ModelV2) + def value_function(self): + return torch.reshape(self._value_out, [-1]) diff --git a/rllib/examples/models/parametric_actions_model.py b/rllib/examples/models/parametric_actions_model.py new file mode 100644 index 000000000..15918d28b --- /dev/null +++ b/rllib/examples/models/parametric_actions_model.py @@ -0,0 +1,111 @@ +from gym.spaces import Box + +from ray.rllib.agents.dqn.distributional_q_tf_model import \ + DistributionalQTFModel +from ray.rllib.agents.dqn.dqn_torch_model import \ + DQNTorchModel +from ray.rllib.models.tf.fcnet_v2 import FullyConnectedNetwork +from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFC +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.numpy import LARGE_INTEGER + +tf = try_import_tf() +torch, nn = try_import_torch() + + +class ParametricActionsModel(DistributionalQTFModel): + """Parametric action model that handles the dot product and masking. + + This assumes the outputs are logits for a single Categorical action dist. + Getting this to work with a more complex output (e.g., if the action space + is a tuple of several distributions) is also possible but left as an + exercise to the reader. + """ + + def __init__(self, + obs_space, + action_space, + num_outputs, + model_config, + name, + true_obs_shape=(4, ), + action_embed_size=2, + **kw): + super(ParametricActionsModel, self).__init__( + obs_space, action_space, num_outputs, model_config, name, **kw) + self.action_embed_model = FullyConnectedNetwork( + Box(-1, 1, shape=true_obs_shape), action_space, action_embed_size, + model_config, name + "_action_embed") + self.register_variables(self.action_embed_model.variables()) + + def forward(self, input_dict, state, seq_lens): + # Extract the available actions tensor from the observation. + avail_actions = input_dict["obs"]["avail_actions"] + action_mask = input_dict["obs"]["action_mask"] + + # Compute the predicted action embedding + action_embed, _ = self.action_embed_model({ + "obs": input_dict["obs"]["cart"] + }) + + # Expand the model output to [BATCH, 1, EMBED_SIZE]. Note that the + # avail actions tensor is of shape [BATCH, MAX_ACTIONS, EMBED_SIZE]. + intent_vector = tf.expand_dims(action_embed, 1) + + # Batch dot product => shape of logits is [BATCH, MAX_ACTIONS]. + action_logits = tf.reduce_sum(avail_actions * intent_vector, axis=2) + + # Mask out invalid actions (use tf.float32.min for stability) + inf_mask = tf.maximum(tf.log(action_mask), tf.float32.min) + return action_logits + inf_mask, state + + def value_function(self): + return self.action_embed_model.value_function() + + +class TorchParametricActionsModel(DQNTorchModel, nn.Module): + """PyTorch version of above ParametricActionsModel.""" + + def __init__(self, + obs_space, + action_space, + num_outputs, + model_config, + name, + true_obs_shape=(4, ), + action_embed_size=2, + **kw): + nn.Module.__init__(self) + DQNTorchModel.__init__(self, obs_space, action_space, num_outputs, + model_config, name, **kw) + + self.action_embed_model = TorchFC( + Box(-1, 1, shape=true_obs_shape), action_space, action_embed_size, + model_config, name + "_action_embed") + + def forward(self, input_dict, state, seq_lens): + # Extract the available actions tensor from the observation. + avail_actions = input_dict["obs"]["avail_actions"] + action_mask = input_dict["obs"]["action_mask"] + + # Compute the predicted action embedding + action_embed, _ = self.action_embed_model({ + "obs": input_dict["obs"]["cart"] + }) + + # Expand the model output to [BATCH, 1, EMBED_SIZE]. Note that the + # avail actions tensor is of shape [BATCH, MAX_ACTIONS, EMBED_SIZE]. + intent_vector = torch.unsqueeze(action_embed, 1) + + # Batch dot product => shape of logits is [BATCH, MAX_ACTIONS]. + action_logits = torch.sum(avail_actions * intent_vector, dim=2) + + # Mask out invalid actions (use -LARGE_INTEGER to tag invalid). + # These are then recognized by the EpsilonGreedy exploration component + # as invalid actions that are not to be chosen. + inf_mask = torch.clamp( + torch.log(action_mask), -float(LARGE_INTEGER), float("inf")) + return action_logits + inf_mask, state + + def value_function(self): + return self.action_embed_model.value_function() diff --git a/rllib/examples/models/rnn_model.py b/rllib/examples/models/rnn_model.py new file mode 100644 index 000000000..963623ae2 --- /dev/null +++ b/rllib/examples/models/rnn_model.py @@ -0,0 +1,138 @@ +import numpy as np + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.preprocessors import get_preprocessor +from ray.rllib.models.tf.recurrent_tf_modelv2 import RecurrentTFModelV2 +from ray.rllib.models.torch.recurrent_torch_model import RecurrentTorchModel +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf = try_import_tf() +torch, nn = try_import_torch() + + +class RNNModel(RecurrentTFModelV2): + """Example of using the Keras functional API to define a RNN model.""" + + def __init__(self, + obs_space, + action_space, + num_outputs, + model_config, + name, + hiddens_size=256, + cell_size=64): + super(RNNModel, self).__init__(obs_space, action_space, num_outputs, + model_config, name) + self.cell_size = cell_size + + # Define input layers + input_layer = tf.keras.layers.Input( + shape=(None, obs_space.shape[0]), name="inputs") + state_in_h = tf.keras.layers.Input(shape=(cell_size, ), name="h") + state_in_c = tf.keras.layers.Input(shape=(cell_size, ), name="c") + seq_in = tf.keras.layers.Input(shape=(), name="seq_in", dtype=tf.int32) + + # Preprocess observation with a hidden layer and send to LSTM cell + dense1 = tf.keras.layers.Dense( + hiddens_size, activation=tf.nn.relu, name="dense1")(input_layer) + lstm_out, state_h, state_c = tf.keras.layers.LSTM( + cell_size, return_sequences=True, return_state=True, name="lstm")( + inputs=dense1, + mask=tf.sequence_mask(seq_in), + initial_state=[state_in_h, state_in_c]) + + # Postprocess LSTM output with another hidden layer and compute values + logits = tf.keras.layers.Dense( + self.num_outputs, + activation=tf.keras.activations.linear, + name="logits")(lstm_out) + values = tf.keras.layers.Dense( + 1, activation=None, name="values")(lstm_out) + + # Create the RNN model + self.rnn_model = tf.keras.Model( + inputs=[input_layer, seq_in, state_in_h, state_in_c], + outputs=[logits, values, state_h, state_c]) + self.register_variables(self.rnn_model.variables) + self.rnn_model.summary() + + @override(RecurrentTFModelV2) + def forward_rnn(self, inputs, state, seq_lens): + model_out, self._value_out, h, c = self.rnn_model([inputs, seq_lens] + + state) + return model_out, [h, c] + + @override(ModelV2) + def get_initial_state(self): + return [ + np.zeros(self.cell_size, np.float32), + np.zeros(self.cell_size, np.float32), + ] + + @override(ModelV2) + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class TorchRNNModel(RecurrentTorchModel): + def __init__(self, + obs_space, + action_space, + num_outputs, + model_config, + name, + fc_size=64, + lstm_state_size=256): + super().__init__(obs_space, action_space, num_outputs, model_config, + name) + + self.obs_size = get_preprocessor(obs_space)(obs_space).size + self.fc_size = fc_size + self.lstm_state_size = lstm_state_size + + # Build the Module from fc + LSTM + 2xfc (action + value outs). + self.fc1 = nn.Linear(self.obs_size, self.fc_size) + self.lstm = nn.LSTM( + self.fc_size, self.lstm_state_size, batch_first=True) + self.action_branch = nn.Linear(self.lstm_state_size, num_outputs) + self.value_branch = nn.Linear(self.lstm_state_size, 1) + # Store the value output to save an extra forward pass. + self._cur_value = None + + @override(ModelV2) + def get_initial_state(self): + # Place hidden states on same device as model. + h = [ + self.fc1.weight.new(1, self.lstm_state_size).zero_().squeeze(0), + self.fc1.weight.new(1, self.lstm_state_size).zero_().squeeze(0) + ] + return h + + @override(ModelV2) + def value_function(self): + assert self._cur_value is not None, "must call forward() first" + return self._cur_value + + @override(RecurrentTorchModel) + def forward_rnn(self, inputs, state, seq_lens): + """Feeds `inputs` (B x T x ..) through the Gru Unit. + + Returns the resulting outputs as a sequence (B x T x ...). + Values are stored in self._cur_value in simple (B) shape (where B + contains both the B and T dims!). + + Returns: + NN Outputs (B x T x ...) as sequence. + The state batches as a List of two items (c- and h-states). + """ + x = nn.functional.relu(self.fc1(inputs)) + lstm_out = self.lstm( + x, [torch.unsqueeze(state[0], 0), + torch.unsqueeze(state[1], 0)]) + action_out = self.action_branch(lstm_out[0]) + self._cur_value = torch.reshape(self.value_branch(lstm_out[0]), [-1]) + return action_out, [ + torch.squeeze(lstm_out[1][0], 0), + torch.squeeze(lstm_out[1][1], 0) + ] diff --git a/rllib/examples/models/shared_weights_model.py b/rllib/examples/models/shared_weights_model.py new file mode 100644 index 000000000..1e68902f9 --- /dev/null +++ b/rllib/examples/models/shared_weights_model.py @@ -0,0 +1,124 @@ +import numpy as np + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils import try_import_tf, try_import_torch + +tf = try_import_tf() +torch, nn = try_import_torch() + + +class SharedWeightsModel1(TFModelV2): + """Example of weight sharing between two different TFModelV2s. + + Here, we share the variables defined in the 'shared' variable scope + by entering it explicitly with tf.AUTO_REUSE. This creates the + variables for the 'fc1' layer in a global scope called 'shared' + (outside of the Policy's normal variable scope). + """ + + def __init__(self, observation_space, action_space, num_outputs, + model_config, name): + super().__init__(observation_space, action_space, num_outputs, + model_config, name) + + inputs = tf.keras.layers.Input(observation_space.shape) + with tf.variable_scope( + tf.VariableScope(tf.AUTO_REUSE, "shared"), + reuse=tf.AUTO_REUSE, + auxiliary_name_scope=False): + last_layer = tf.keras.layers.Dense( + units=64, activation=tf.nn.relu, name="fc1")(inputs) + output = tf.keras.layers.Dense( + units=num_outputs, activation=None, name="fc_out")(last_layer) + vf = tf.keras.layers.Dense( + units=1, activation=None, name="value_out")(last_layer) + self.base_model = tf.keras.models.Model(inputs, [output, vf]) + self.register_variables(self.base_model.variables) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + out, self._value_out = self.base_model(input_dict["obs"]) + return out, [] + + @override(ModelV2) + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +class SharedWeightsModel2(TFModelV2): + """The "other" TFModelV2 using the same shared space as the one above.""" + + def __init__(self, observation_space, action_space, num_outputs, + model_config, name): + super().__init__(observation_space, action_space, num_outputs, + model_config, name) + + inputs = tf.keras.layers.Input(observation_space.shape) + + # Weights shared with SharedWeightsModel1. + with tf.variable_scope( + tf.VariableScope(tf.AUTO_REUSE, "shared"), + reuse=tf.AUTO_REUSE, + auxiliary_name_scope=False): + last_layer = tf.keras.layers.Dense( + units=64, activation=tf.nn.relu, name="fc1")(inputs) + output = tf.keras.layers.Dense( + units=num_outputs, activation=None, name="fc_out")(last_layer) + vf = tf.keras.layers.Dense( + units=1, activation=None, name="value_out")(last_layer) + self.base_model = tf.keras.models.Model(inputs, [output, vf]) + self.register_variables(self.base_model.variables) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + out, self._value_out = self.base_model(input_dict["obs"]) + return out, [] + + @override(ModelV2) + def value_function(self): + return tf.reshape(self._value_out, [-1]) + + +TORCH_GLOBAL_SHARED_LAYER = None +if torch: + TORCH_GLOBAL_SHARED_LAYER = SlimFC(32, 32) + + +class TorchSharedWeightsModel(TorchModelV2, nn.Module): + """Example of weight sharing between two different TorchModelV2s. + + The shared (single) layer is simply defined outside of the two Models, + then used by both Models in their forward pass. + """ + + def __init__(self, observation_space, action_space, num_outputs, + model_config, name): + TorchModelV2.__init__(self, observation_space, action_space, + num_outputs, model_config, name) + nn.Module.__init__(self) + + # Non-shared initial layer. + self.first_layer = SlimFC( + int(np.product(observation_space.shape)), + 32, + activation_fn=nn.ReLU) + + # Non-shared final layer. + self.last_layer = SlimFC(32, self.num_outputs, activation_fn=nn.ReLU) + self.vf = SlimFC(32, 1, activation_fn=None) + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + out = self.first_layer(input_dict["obs"]) + out = TORCH_GLOBAL_SHARED_LAYER(out) + model_out = self.last_layer(out) + self._value_out = self.vf(out) + return model_out, [] + + @override(ModelV2) + def value_function(self): + return torch.reshape(self._value_out, [-1]) diff --git a/rllib/examples/policy/__init__.py b/rllib/examples/policy/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/rllib/examples/policy/random_policy.py b/rllib/examples/policy/random_policy.py new file mode 100644 index 000000000..4cd2d0206 --- /dev/null +++ b/rllib/examples/policy/random_policy.py @@ -0,0 +1,37 @@ +import numpy as np +import random + +from ray.rllib.policy.policy import Policy +from ray.rllib.utils.annotations import override + + +class RandomPolicy(Policy): + """Hand-coded policy that returns random actions.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + @override(Policy) + def compute_actions(self, + obs_batch, + state_batches=None, + prev_action_batch=None, + prev_reward_batch=None, + **kwargs): + # Alternatively, a numpy array would work here as well. + # e.g.: np.array([random.choice([0, 1])] * len(obs_batch)) + return [self.action_space.sample() for _ in obs_batch], [], {} + + @override(Policy) + def learn_on_batch(self, samples): + """No learning.""" + return {} + + @override(Policy) + def compute_log_likelihoods(self, + actions, + obs_batch, + state_batches=None, + prev_action_batch=None, + prev_reward_batch=None): + return np.array([random.random()] * len(obs_batch)) diff --git a/rllib/examples/policy/rock_paper_scissors_dummies.py b/rllib/examples/policy/rock_paper_scissors_dummies.py new file mode 100644 index 000000000..2eb051b09 --- /dev/null +++ b/rllib/examples/policy/rock_paper_scissors_dummies.py @@ -0,0 +1,65 @@ +import random + +from ray.rllib.examples.env.rock_paper_scissors import RockPaperScissors +from ray.rllib.policy.policy import Policy + + +class AlwaysSameHeuristic(Policy): + """Pick a random move and stick with it for the entire episode.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.exploration = self._create_exploration() + + def get_initial_state(self): + return [ + random.choice([ + RockPaperScissors.ROCK, RockPaperScissors.PAPER, + RockPaperScissors.SCISSORS + ]) + ] + + def compute_actions(self, + obs_batch, + state_batches=None, + prev_action_batch=None, + prev_reward_batch=None, + info_batch=None, + episodes=None, + **kwargs): + return state_batches[0], state_batches, {} + + +class BeatLastHeuristic(Policy): + """Play the move that would beat the last move of the opponent.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.exploration = self._create_exploration() + + def compute_actions(self, + obs_batch, + state_batches=None, + prev_action_batch=None, + prev_reward_batch=None, + info_batch=None, + episodes=None, + **kwargs): + def successor(x): + if x[RockPaperScissors.ROCK] == 1: + return RockPaperScissors.PAPER + elif x[RockPaperScissors.PAPER] == 1: + return RockPaperScissors.SCISSORS + elif x[RockPaperScissors.SCISSORS] == 1: + return RockPaperScissors.ROCK + + return [successor(x) for x in obs_batch], [], {} + + def learn_on_batch(self, samples): + pass + + def get_weights(self): + pass + + def set_weights(self, weights): + pass diff --git a/rllib/models/catalog.py b/rllib/models/catalog.py index c00496d79..c6bdc507e 100644 --- a/rllib/models/catalog.py +++ b/rllib/models/catalog.py @@ -169,13 +169,13 @@ class ModelCatalog: lambda s: ModelCatalog.get_action_dist( s, config, framework=framework), flat_action_space) child_dists = [e[0] for e in child_dists_and_in_lens] - input_lens = [e[1] for e in child_dists_and_in_lens] + input_lens = [int(e[1]) for e in child_dists_and_in_lens] return partial( (TorchMultiActionDistribution if framework == "torch" else MultiActionDistribution), action_space=action_space, child_distributions=child_dists, - input_lens=input_lens), sum(input_lens) + input_lens=input_lens), int(sum(input_lens)) # Simplex -> Dirichlet. elif isinstance(action_space, Simplex): if framework == "torch": @@ -281,14 +281,16 @@ class ModelCatalog: model_cls = _global_registry.get(RLLIB_MODEL, model_config["custom_model"]) if issubclass(model_cls, ModelV2): + + logger.info("Wrapping {} as {}".format(model_cls, + model_interface)) + model_cls = ModelCatalog._wrap_if_needed( + model_cls, model_interface) + if framework == "tf": - logger.info("Wrapping {} as {}".format( - model_cls, model_interface)) - model_cls = ModelCatalog._wrap_if_needed( - model_cls, model_interface) + # Track and warn if vars were created but not registered. created = set() - # Track and warn if vars were created but not registered def track_var_creation(next_creator, **kw): v = next_creator(**kw) created.add(v) @@ -312,10 +314,13 @@ class ModelCatalog: "question?".format(not_registered, instance, registered)) else: - # no variable tracking + # PyTorch automatically tracks nn.Modules inside the parent + # nn.Module's constructor. + # TODO(sven): Do this for TF as well. instance = model_cls(obs_space, action_space, num_outputs, model_config, name, **model_kwargs) return instance + elif tf.executing_eagerly(): raise ValueError( "Eager execution requires a TFModelV2 model to be " diff --git a/rllib/models/torch/fcnet.py b/rllib/models/torch/fcnet.py index f44a3b9a8..21e5bb255 100644 --- a/rllib/models/torch/fcnet.py +++ b/rllib/models/torch/fcnet.py @@ -88,7 +88,7 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module): @override(TorchModelV2) def forward(self, input_dict, state, seq_lens): - obs = input_dict["obs_flat"] + obs = input_dict["obs_flat"].float() features = self._hidden_layers(obs.reshape(obs.shape[0], -1)) logits = self._logits(features) if self._logits else features self._cur_value = self._value_branch(features).squeeze(1) diff --git a/rllib/models/torch/recurrent_torch_model.py b/rllib/models/torch/recurrent_torch_model.py index 0d99e0bfb..dce8eb95d 100644 --- a/rllib/models/torch/recurrent_torch_model.py +++ b/rllib/models/torch/recurrent_torch_model.py @@ -31,7 +31,7 @@ class RecurrentTorchModel(TorchModelV2, nn.Module): @override(ModelV2) def get_initial_state(self): - # make hidden states on same device as model + # Place hidden states on same device as model. h = [self.fc1.weight.new( 1, self.rnn_hidden_dim).zero_().squeeze(0)] return h diff --git a/rllib/policy/dynamic_tf_policy.py b/rllib/policy/dynamic_tf_policy.py index bf5f8ca70..e301c2084 100644 --- a/rllib/policy/dynamic_tf_policy.py +++ b/rllib/policy/dynamic_tf_policy.py @@ -156,11 +156,15 @@ class DynamicTFPolicy(TFPolicy): self.model = make_model(self, obs_space, action_space, config) else: self.model = ModelCatalog.get_model_v2( - obs_space, - action_space, - logit_dim, - self.config["model"], + obs_space=obs_space, + action_space=action_space, + num_outputs=logit_dim, + model_config=self.config["model"], framework="tf") + # NOTE: Adding below line will break existing custom models + # that do not expect extra options in **kwargs but rather in + # model_config["custom_options"]. + # **self.config["model"].get("custom_options", {})) # Create the Exploration object to use for this Policy. self.exploration = self._create_exploration() diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py index fe29753f3..8056978d9 100644 --- a/rllib/policy/torch_policy.py +++ b/rllib/policy/torch_policy.py @@ -232,7 +232,7 @@ class TorchPolicy(Policy): loss_out = force_list( self._loss(self, self.model, self.dist_class, train_batch)) assert len(loss_out) == len(self._optimizers) - assert not any(np.isnan(l.detach().numpy()) for l in loss_out) + # assert not any(np.isnan(l.detach().numpy()) for l in loss_out) # Loop through all optimizers. grad_info = {"allreduce_latency": 0.0} diff --git a/rllib/rollout.py b/rllib/rollout.py index 841ac2e21..3df0233e1 100755 --- a/rllib/rollout.py +++ b/rllib/rollout.py @@ -32,6 +32,10 @@ Example Usage via executable: # Note: if you use any custom models or envs, register them here first, e.g.: # +# from ray.rllib.examples.env.parametric_actions_cartpole import \ +# ParametricActionsCartPole +# from ray.rllib.examples.model.parametric_actions_model import \ +# ParametricActionsModel # ModelCatalog.register_custom_model("pa_model", ParametricActionsModel) # register_env("pa_cartpole", lambda _: ParametricActionsCartPole(10)) diff --git a/rllib/utils/exploration/epsilon_greedy.py b/rllib/utils/exploration/epsilon_greedy.py index 29d5ae8fa..26abd766a 100644 --- a/rllib/utils/exploration/epsilon_greedy.py +++ b/rllib/utils/exploration/epsilon_greedy.py @@ -6,6 +6,7 @@ from ray.rllib.utils.exploration.exploration import Exploration, TensorType from ray.rllib.utils.framework import try_import_tf, try_import_torch, \ get_variable from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.numpy import LARGE_INTEGER from ray.rllib.utils.schedules import Schedule, PiecewiseSchedule tf = try_import_tf() @@ -136,9 +137,8 @@ class EpsilonGreedy(Exploration): # Mask out actions, whose Q-values are -inf, so that we don't # even consider them for exploration. random_valid_action_logits = torch.where( - q_values == float("-inf"), - torch.ones_like(q_values) * float("-inf"), - torch.ones_like(q_values)) + q_values == -float(LARGE_INTEGER), + torch.ones_like(q_values) * 0.0, torch.ones_like(q_values)) # A random action. random_actions = torch.squeeze( torch.multinomial(random_valid_action_logits, 1), axis=1) diff --git a/rllib/utils/exploration/exploration.py b/rllib/utils/exploration/exploration.py index 9a1145526..ca650e82e 100644 --- a/rllib/utils/exploration/exploration.py +++ b/rllib/utils/exploration/exploration.py @@ -39,8 +39,11 @@ class Exploration: self.framework = check_framework(framework) # The device on which the Model has been placed. # This Exploration will be on the same device. - self.device = None if not isinstance(self.model, nn.Module) else \ - next(self.model.parameters()).device + self.device = None + if isinstance(self.model, nn.Module): + params = list(self.model.parameters()) + if params: + self.device = params[0].device @DeveloperAPI def before_compute_actions(self, diff --git a/rllib/utils/test_utils.py b/rllib/utils/test_utils.py index 600040d91..aa07d90e4 100644 --- a/rllib/utils/test_utils.py +++ b/rllib/utils/test_utils.py @@ -225,3 +225,21 @@ def check(x, y, decimals=5, atol=None, rtol=None, false=False): if false is True: assert False, \ "ERROR: x ({}) is the same as y ({})!".format(x, y) + + +def check_learning_achieved(tune_results, min_reward): + """Throws an error if `min_reward` is not reached within tune_results. + + Checks the last iteration found in tune_results for its + "episode_reward_mean" value and compares it to `min_reward`. + + Args: + tune_results: The tune.run returned results object. + min_reward (float): The min reward that must be reached. + + Throws: + ValueError: If `min_reward` not reached. + """ + if tune_results.trials[0].last_result["episode_reward_mean"] < min_reward: + raise ValueError("`stop-reward` of {} not reached!".format(min_reward)) + print("ok")