diff --git a/rllib/agents/ppo/tests/test_ppo.py b/rllib/agents/ppo/tests/test_ppo.py index 14530fcd0..a8bc87ccc 100644 --- a/rllib/agents/ppo/tests/test_ppo.py +++ b/rllib/agents/ppo/tests/test_ppo.py @@ -1,3 +1,4 @@ +import copy import numpy as np import unittest @@ -19,6 +20,22 @@ from ray.rllib.utils.test_utils import check, framework_iterator, \ tf = try_import_tf() +# Fake CartPole episode of n time steps. +FAKE_BATCH = { + SampleBatch.CUR_OBS: np.array( + [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1.0, 1.1, 1.2]], + dtype=np.float32), + SampleBatch.ACTIONS: np.array([0, 1, 1]), + SampleBatch.PREV_ACTIONS: np.array([0, 1, 1]), + SampleBatch.REWARDS: np.array([1.0, -1.0, .5], dtype=np.float32), + SampleBatch.PREV_REWARDS: np.array([1.0, -1.0, .5], dtype=np.float32), + SampleBatch.DONES: np.array([False, False, True]), + SampleBatch.VF_PREDS: np.array([0.5, 0.6, 0.7], dtype=np.float32), + SampleBatch.ACTION_DIST_INPUTS: np.array( + [[-2., 0.5], [-3., -0.3], [-0.1, 2.5]], dtype=np.float32), + SampleBatch.ACTION_LOGP: np.array([-0.5, -0.1, -0.2], dtype=np.float32), +} + class TestPPO(unittest.TestCase): @classmethod @@ -31,7 +48,7 @@ class TestPPO(unittest.TestCase): def test_ppo_compilation(self): """Test whether a PPOTrainer can be built with both frameworks.""" - config = ppo.DEFAULT_CONFIG.copy() + config = copy.deepcopy(ppo.DEFAULT_CONFIG) config["num_workers"] = 0 # Run locally. num_iterations = 2 @@ -43,7 +60,7 @@ class TestPPO(unittest.TestCase): def test_ppo_fake_multi_gpu_learning(self): """Test whether PPOTrainer can learn CartPole w/ faked multi-GPU.""" - config = ppo.DEFAULT_CONFIG.copy() + config = copy.deepcopy(ppo.DEFAULT_CONFIG) # Fake GPU setup. config["num_gpus"] = 2 config["_fake_gpus"] = True @@ -70,7 +87,7 @@ class TestPPO(unittest.TestCase): def test_ppo_exploration_setup(self): """Tests, whether PPO runs with different exploration setups.""" - config = ppo.DEFAULT_CONFIG.copy() + config = copy.deepcopy(ppo.DEFAULT_CONFIG) config["num_workers"] = 0 # Run locally. config["env_config"] = {"is_slippery": False, "map_name": "4x4"} obs = np.array(0) @@ -107,46 +124,91 @@ class TestPPO(unittest.TestCase): prev_reward=np.array(1.0))) check(np.mean(actions), 1.5, atol=0.2) + def test_ppo_free_log_std(self): + """Tests the free log std option works.""" + config = copy.deepcopy(ppo.DEFAULT_CONFIG) + config["num_workers"] = 0 # Run locally. + config["gamma"] = 0.99 + config["model"]["fcnet_hiddens"] = [10] + config["model"]["fcnet_activation"] = "linear" + config["model"]["free_log_std"] = True + config["vf_share_layers"] = True + + for fw, sess in framework_iterator(config, session=True): + trainer = ppo.PPOTrainer(config=config, env="CartPole-v0") + policy = trainer.get_policy() + + # Check the free log std var is created. + if fw == "torch": + matching = [ + v for (n, v) in policy.model.named_parameters() + if "log_std" in n + ] + else: + matching = [ + v for v in policy.model.trainable_variables() + if "log_std" in str(v) + ] + assert len(matching) == 1, matching + log_std_var = matching[0] + + def get_value(): + if fw == "tf": + return policy.get_session().run(log_std_var)[0] + elif fw == "torch": + return log_std_var.detach().numpy()[0] + else: + return log_std_var.numpy()[0] + + # Check the variable is initially zero. + init_std = get_value() + assert init_std == 0.0, init_std + + if fw == "tf" or fw == "eager": + batch = postprocess_ppo_gae_tf(policy, FAKE_BATCH) + else: + batch = postprocess_ppo_gae_torch(policy, FAKE_BATCH) + batch = policy._lazy_tensor_dict(batch) + policy.learn_on_batch(batch) + + # Check the variable is updated. + post_std = get_value() + assert post_std != 0.0, post_std + def test_ppo_loss_function(self): """Tests the PPO loss function math.""" - config = ppo.DEFAULT_CONFIG.copy() + config = copy.deepcopy(ppo.DEFAULT_CONFIG) config["num_workers"] = 0 # Run locally. config["gamma"] = 0.99 config["model"]["fcnet_hiddens"] = [10] config["model"]["fcnet_activation"] = "linear" config["vf_share_layers"] = True - # Fake CartPole episode of n time steps. - train_batch = { - SampleBatch.CUR_OBS: np.array( - [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], - [0.9, 1.0, 1.1, 1.2]], - dtype=np.float32), - SampleBatch.ACTIONS: np.array([0, 1, 1]), - SampleBatch.PREV_ACTIONS: np.array([0, 1, 1]), - SampleBatch.REWARDS: np.array([1.0, -1.0, .5], dtype=np.float32), - SampleBatch.PREV_REWARDS: np.array( - [1.0, -1.0, .5], dtype=np.float32), - SampleBatch.DONES: np.array([False, False, True]), - SampleBatch.VF_PREDS: np.array([0.5, 0.6, 0.7], dtype=np.float32), - SampleBatch.ACTION_DIST_INPUTS: np.array( - [[-2., 0.5], [-3., -0.3], [-0.1, 2.5]], dtype=np.float32), - SampleBatch.ACTION_LOGP: np.array( - [-0.5, -0.1, -0.2], dtype=np.float32), - } - for fw, sess in framework_iterator(config, session=True): trainer = ppo.PPOTrainer(config=config, env="CartPole-v0") policy = trainer.get_policy() + # Check no free log std var by default. + if fw == "torch": + matching = [ + v for (n, v) in policy.model.named_parameters() + if "log_std" in n + ] + else: + matching = [ + v for v in policy.model.trainable_variables() + if "log_std" in str(v) + ] + assert len(matching) == 0, matching + # Post-process (calculate simple (non-GAE) advantages) and attach # to train_batch dict. # A = [0.99^2 * 0.5 + 0.99 * -1.0 + 1.0, 0.99 * 0.5 - 1.0, 0.5] = # [0.50005, -0.505, 0.5] if fw == "tf" or fw == "eager": - train_batch = postprocess_ppo_gae_tf(policy, train_batch) + train_batch = postprocess_ppo_gae_tf(policy, FAKE_BATCH) else: - train_batch = postprocess_ppo_gae_torch(policy, train_batch) + train_batch = postprocess_ppo_gae_torch(policy, FAKE_BATCH) train_batch = policy._lazy_tensor_dict(train_batch) # Check Advantage values. diff --git a/rllib/models/catalog.py b/rllib/models/catalog.py index c6bdc507e..b646fc452 100644 --- a/rllib/models/catalog.py +++ b/rllib/models/catalog.py @@ -44,7 +44,9 @@ MODEL_DEFAULTS = { "fcnet_activation": "tanh", # Number of hidden layers for fully connected net "fcnet_hiddens": [256, 256], - # For control envs, documented in ray.rllib.models.Model + # For DiagGaussian action distributions, make the second half of the model + # outputs floating bias variables instead of state-dependent. This only + # has an effect is using the default fully connected net. "free_log_std": False, # Whether to skip the final linear layer used to resize the hidden layer # outputs to size `num_outputs`. If True, then the last hidden layer diff --git a/rllib/models/tf/fcnet_v2.py b/rllib/models/tf/fcnet_v2.py index 303016f27..745639a98 100644 --- a/rllib/models/tf/fcnet_v2.py +++ b/rllib/models/tf/fcnet_v2.py @@ -19,8 +19,19 @@ class FullyConnectedNetwork(TFModelV2): hiddens = model_config.get("fcnet_hiddens", []) no_final_linear = model_config.get("no_final_linear") vf_share_layers = model_config.get("vf_share_layers") + free_log_std = model_config.get("free_log_std") - # we are using obs_flat, so take the flattened shape as input + # Maybe generate free-floating bias variables for the second half of + # the outputs. + if free_log_std: + assert num_outputs % 2 == 0, ( + "num_outputs must be divisible by two", num_outputs) + num_outputs = num_outputs // 2 + self.log_std_var = tf.Variable( + [0.0] * num_outputs, dtype=tf.float32, name="log_std") + self.register_variables([self.log_std_var]) + + # We are using obs_flat, so take the flattened shape as input. inputs = tf.keras.layers.Input( shape=(np.product(obs_space.shape), ), name="observations") last_layer = layer_out = inputs @@ -37,9 +48,9 @@ class FullyConnectedNetwork(TFModelV2): # The last layer is adjusted to be of size num_outputs, but it's a # layer with activation. - if no_final_linear and self.num_outputs: + if no_final_linear and num_outputs: layer_out = tf.keras.layers.Dense( - self.num_outputs, + num_outputs, name="fc_out", activation=activation, kernel_initializer=normc_initializer(1.0))(last_layer) @@ -52,18 +63,28 @@ class FullyConnectedNetwork(TFModelV2): name="fc_{}".format(i), activation=activation, kernel_initializer=normc_initializer(1.0))(last_layer) - if self.num_outputs: + if num_outputs: layer_out = tf.keras.layers.Dense( - self.num_outputs, + num_outputs, name="fc_out", activation=None, kernel_initializer=normc_initializer(0.01))(last_layer) - # Adjust self.num_outputs to be the number of nodes in the last - # layer. + # Adjust num_outputs to be the number of nodes in the last layer. else: self.num_outputs = ( [np.product(obs_space.shape)] + hiddens[-1:-1])[-1] + # Concat the log std vars to the end of the state-dependent means. + if free_log_std: + + def tiled_log_std(x): + return tf.tile( + tf.expand_dims(self.log_std_var, 0), [tf.shape(x)[0], 1]) + + log_std_out = tf.keras.layers.Lambda(tiled_log_std)(inputs) + layer_out = tf.keras.layers.Concatenate(axis=1)( + [layer_out, log_std_out]) + if not vf_share_layers: # build a parallel set of hidden layers for the value net last_layer = inputs diff --git a/rllib/models/torch/fcnet.py b/rllib/models/torch/fcnet.py index 21e5bb255..4065309bb 100644 --- a/rllib/models/torch/fcnet.py +++ b/rllib/models/torch/fcnet.py @@ -2,12 +2,13 @@ import logging import numpy as np from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 -from ray.rllib.models.torch.misc import SlimFC, normc_initializer +from ray.rllib.models.torch.misc import SlimFC, AppendBiasLayer, \ + normc_initializer from ray.rllib.utils.annotations import override from ray.rllib.utils.framework import get_activation_fn from ray.rllib.utils import try_import_torch -_, nn = try_import_torch() +torch, nn = try_import_torch() logger = logging.getLogger(__name__) @@ -25,6 +26,7 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module): model_config.get("fcnet_activation"), framework="torch") hiddens = model_config.get("fcnet_hiddens") no_final_linear = model_config.get("no_final_linear") + self.free_log_std = model_config.get("free_log_std") # TODO(sven): implement case: vf_shared_layers = False. # vf_share_layers = model_config.get("vf_share_layers") @@ -34,6 +36,13 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module): prev_layer_size = int(np.product(obs_space.shape)) self._logits = None + # Maybe generate free-floating bias variables for the second half of + # the outputs. + if self.free_log_std: + assert num_outputs % 2 == 0, ( + "num_outputs must be divisible by two", num_outputs) + num_outputs = num_outputs // 2 + # Create layers 0 to second-last. for size in hiddens[:-1]: layers.append( @@ -46,14 +55,14 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module): # The last layer is adjusted to be of size num_outputs, but it's a # layer with activation. - if no_final_linear and self.num_outputs: + if no_final_linear and num_outputs: layers.append( SlimFC( in_size=prev_layer_size, - out_size=self.num_outputs, + out_size=num_outputs, initializer=normc_initializer(1.0), activation_fn=activation)) - prev_layer_size = self.num_outputs + prev_layer_size = num_outputs # Finish the layers with the provided sizes (`hiddens`), plus - # iff num_outputs > 0 - a last linear layer of size num_outputs. else: @@ -65,16 +74,20 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module): initializer=normc_initializer(1.0), activation_fn=activation)) prev_layer_size = hiddens[-1] - if self.num_outputs: + if num_outputs: self._logits = SlimFC( in_size=prev_layer_size, - out_size=self.num_outputs, + out_size=num_outputs, initializer=normc_initializer(0.01), activation_fn=None) else: self.num_outputs = ( [np.product(obs_space.shape)] + hiddens[-1:-1])[-1] + # Layer to add the log std vars to the state-dependent means. + if self.free_log_std: + self._append_free_log_std = AppendBiasLayer(num_outputs) + self._hidden_layers = nn.Sequential(*layers) # TODO(sven): Implement non-shared value branch. @@ -91,6 +104,8 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module): 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 + if self.free_log_std: + logits = self._append_free_log_std(logits) self._cur_value = self._value_branch(features).squeeze(1) return logits, state diff --git a/rllib/models/torch/misc.py b/rllib/models/torch/misc.py index d86368a80..71b223187 100644 --- a/rllib/models/torch/misc.py +++ b/rllib/models/torch/misc.py @@ -106,3 +106,18 @@ class SlimFC(nn.Module): def forward(self, x): return self._model(x) + + +class AppendBiasLayer(nn.Module): + """Simple bias appending layer for free_log_std.""" + + def __init__(self, num_bias_vars): + super().__init__() + self.log_std = torch.nn.Parameter( + torch.as_tensor([0.0] * num_bias_vars)) + self.register_parameter("log_std", self.log_std) + + def forward(self, x): + out = torch.cat( + [x, self.log_std.unsqueeze(0).repeat([len(x), 1])], axis=1) + return out