diff --git a/rllib/BUILD b/rllib/BUILD index cc500ce19..de99dd58b 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -1171,7 +1171,7 @@ py_test( name = "test_env_with_subprocess", main = "tests/test_env_with_subprocess.py", tags = ["tests_dir", "tests_dir_E"], - size = "small", + size = "medium", srcs = ["tests/test_env_with_subprocess.py"] ) diff --git a/rllib/agents/dqn/distributional_q_tf_model.py b/rllib/agents/dqn/distributional_q_tf_model.py index b04ab191f..61045ae1c 100644 --- a/rllib/agents/dqn/distributional_q_tf_model.py +++ b/rllib/agents/dqn/distributional_q_tf_model.py @@ -49,12 +49,12 @@ class DistributionalQTFModel(TFModelV2): for DDQN. If True, Q-values are calculated as: Q = (A - mean[A]) + V. If False, raw NN output is interpreted as Q-values. - num_atoms (int): if >1, enables distributional DQN - use_noisy (bool): use noisy nets - v_min (float): min value support for distributional DQN - v_max (float): max value support for distributional DQN - sigma0 (float): initial value of noisy nets - add_layer_norm (bool): Add a LayerNorm after each layer.. + num_atoms (int): If >1, enables distributional DQN. + use_noisy (bool): Use noisy nets. + v_min (float): Min value support for distributional DQN. + v_max (float): Max value support for distributional DQN. + sigma0 (float): Initial value of noisy layers. + add_layer_norm (bool): Enable layer norm (for param noise). Note that the core layers for forward() are not defined here, this only defines the layers for the Q head. Those layers for forward() diff --git a/rllib/agents/dqn/dqn_tf_policy.py b/rllib/agents/dqn/dqn_tf_policy.py index a028031d7..46ec8e28e 100644 --- a/rllib/agents/dqn/dqn_tf_policy.py +++ b/rllib/agents/dqn/dqn_tf_policy.py @@ -231,8 +231,8 @@ def build_q_losses(policy, model, _, train_batch): train_batch[SampleBatch.NEXT_OBS], explore=False) q_tp1_best_using_online_net = tf.argmax(q_tp1_using_online_net, 1) - q_tp1_best_one_hot_selection = tf.one_hot(q_tp1_best_using_online_net, - policy.action_space.n) + q_tp1_best_one_hot_selection = tf.one_hot( + q_tp1_best_using_online_net, policy.action_space.n) q_tp1_best = tf.reduce_sum(q_tp1 * q_tp1_best_one_hot_selection, 1) q_dist_tp1_best = tf.reduce_sum( q_dist_tp1 * tf.expand_dims(q_tp1_best_one_hot_selection, -1), 1) @@ -246,9 +246,9 @@ def build_q_losses(policy, model, _, train_batch): policy.q_loss = QLoss( q_t_selected, q_logits_t_selected, q_tp1_best, q_dist_tp1_best, train_batch[PRIO_WEIGHTS], train_batch[SampleBatch.REWARDS], - tf.cast(train_batch[SampleBatch.DONES], - tf.float32), config["gamma"], config["n_step"], - config["num_atoms"], config["v_min"], config["v_max"]) + tf.cast(train_batch[SampleBatch.DONES], tf.float32), config["gamma"], + config["n_step"], config["num_atoms"], + config["v_min"], config["v_max"]) return policy.q_loss.loss diff --git a/rllib/agents/dqn/dqn_torch_model.py b/rllib/agents/dqn/dqn_torch_model.py index a6978da24..d10c4f013 100644 --- a/rllib/agents/dqn/dqn_torch_model.py +++ b/rllib/agents/dqn/dqn_torch_model.py @@ -1,5 +1,5 @@ -import numpy as np - +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.modules.noisy_layer import NoisyLayer from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.utils.framework import try_import_torch @@ -18,10 +18,13 @@ class DQNTorchModel(TorchModelV2, nn.Module): model_config, name, *, - dueling=False, q_hiddens=(256, ), + dueling=False, dueling_activation="relu", + num_atoms=1, use_noisy=False, + v_min=-10.0, + v_max=10.0, sigma0=0.5, # TODO(sven): Move `add_layer_norm` into ModelCatalog as # generic option, then error if we use ParameterNoise as @@ -31,19 +34,22 @@ class DQNTorchModel(TorchModelV2, nn.Module): """Initialize variables of this model. Extra model kwargs: - dueling (bool): Whether to build the advantage(A)/value(V) heads - for DDQN. If True, Q-values are calculated as: - Q = (A - mean[A]) + V. If False, raw NN output is interpreted - as Q-values. q_hiddens (List[int]): List of layer-sizes after(!) the Advantages(A)/Value(V)-split. Hence, each of the A- and V- branches will have this structure of Dense layers. To define the NN before this A/V-split, use - as always - config["model"]["fcnet_hiddens"]. + dueling (bool): Whether to build the advantage(A)/value(V) heads + for DDQN. If True, Q-values are calculated as: + Q = (A - mean[A]) + V. If False, raw NN output is interpreted + as Q-values. dueling_activation (str): The activation to use for all dueling layers (A- and V-branch). One of "relu", "tanh", "linear". - use_noisy (bool): use noisy nets - sigma0 (float): initial value of noisy nets + num_atoms (int): If >1, enables distributional DQN. + use_noisy (bool): Use noisy layers. + v_min (float): Min value support for distributional DQN. + v_max (float): Max value support for distributional DQN. + sigma0 (float): Initial value of noisy layers. add_layer_norm (bool): Enable layer norm (for param noise). """ nn.Module.__init__(self) @@ -51,6 +57,10 @@ class DQNTorchModel(TorchModelV2, nn.Module): num_outputs, model_config, name) self.dueling = dueling + self.num_atoms = num_atoms + self.v_min = v_min + self.v_max = v_max + self.sigma0 = sigma0 ins = num_outputs advantage_module = nn.Sequential() @@ -58,109 +68,87 @@ class DQNTorchModel(TorchModelV2, nn.Module): # Dueling case: Build the shared (advantages and value) fc-network. for i, n in enumerate(q_hiddens): - advantage_module.add_module("dueling_A_{}".format(i), - nn.Linear(ins, n)) - value_module.add_module("dueling_V_{}".format(i), - nn.Linear(ins, n)) - # Add activations if necessary. - if dueling_activation == "relu": - advantage_module.add_module("dueling_A_act_{}".format(i), - nn.ReLU()) - value_module.add_module("dueling_V_act_{}".format(i), - nn.ReLU()) - elif dueling_activation == "tanh": - advantage_module.add_module("dueling_A_act_{}".format(i), - nn.Tanh()) - value_module.add_module("dueling_V_act_{}".format(i), - nn.Tanh()) - - # Add LayerNorm after each Dense. - if add_layer_norm: - advantage_module.add_module("LayerNorm_A_{}".format(i), - nn.LayerNorm(n)) - value_module.add_module("LayerNorm_V_{}".format(i), - nn.LayerNorm(n)) + if use_noisy: + advantage_module.add_module( + "dueling_A_{}".format(i), + NoisyLayer( + ins, n, sigma0=self.sigma0, + activation=dueling_activation)) + value_module.add_module( + "dueling_V_{}".format(i), + NoisyLayer( + ins, n, sigma0=self.sigma0, + activation=dueling_activation)) + else: + advantage_module.add_module( + "dueling_A_{}".format(i), + SlimFC(ins, n, activation_fn=dueling_activation)) + value_module.add_module( + "dueling_V_{}".format(i), + SlimFC(ins, n, activation_fn=dueling_activation)) + # Add LayerNorm after each Dense. + if add_layer_norm: + advantage_module.add_module( + "LayerNorm_A_{}".format(i), nn.LayerNorm(n)) + value_module.add_module( + "LayerNorm_V_{}".format(i), nn.LayerNorm(n)) ins = n # Actual Advantages layer (nodes=num-actions). - if q_hiddens: - advantage_module.add_module("A", nn.Linear(ins, action_space.n)) + if use_noisy: + advantage_module.add_module("A", NoisyLayer( + ins, + self.action_space.n * self.num_atoms, + sigma0, + activation=None)) + elif q_hiddens: + advantage_module.add_module( + "A", + SlimFC( + ins, action_space.n * self.num_atoms, + activation_fn=None)) self.advantage_module = advantage_module # Value layer (nodes=1). if self.dueling: - value_module.add_module("V", nn.Linear(ins, 1)) + value_module.add_module("V", SlimFC(ins, 1, activation_fn=None)) self.value_module = value_module - def get_advantages_or_q_values(self, model_out): + def get_q_value_distributions(self, model_out): """Returns distributional values for Q(s, a) given a state embedding. Override this in your custom model to customize the Q output head. - Arguments: - model_out (Tensor): embedding from the model layers + Args: + model_out (Tensor): Embedding from the model layers. Returns: (action_scores, logits, dist) if num_atoms == 1, otherwise (action_scores, z, support_logits_per_action, logits, dist) """ + action_scores = self.advantage_module(model_out) - return self.advantage_module(model_out) + if self.num_atoms > 1: + # Distributional Q-learning uses a discrete support z + # to represent the action value distribution + z = torch.range(0.0, self.num_atoms - 1, dtype=torch.float32) + z = self.v_min + \ + z * (self.v_max - self.v_min) / float(self.num_atoms - 1) + + support_logits_per_action = torch.reshape( + action_scores, shape=(-1, self.action_space.n, self.num_atoms)) + support_prob_per_action = nn.functional.softmax( + support_logits_per_action) + action_scores = torch.sum(z * support_prob_per_action, dim=-1) + logits = support_logits_per_action + probs = support_prob_per_action + return action_scores, z, support_logits_per_action, logits, probs + else: + logits = torch.unsqueeze(torch.ones_like(action_scores), -1) + return action_scores, logits, logits def get_state_value(self, model_out): """Returns the state value prediction for the given state embedding.""" return self.value_module(model_out) - - def _noisy_layer(self, action_in, out_size, sigma0, non_linear=True): - """ - a common dense layer: y = w^{T}x + b - a noisy layer: y = (w + \\epsilon_w*\\sigma_w)^{T}x + - (b+\\epsilon_b*\\sigma_b) - where \epsilon are random variables sampled from factorized normal - distributions and \\sigma are trainable variables which are expected to - vanish along the training procedure - """ - in_size = int(action_in.shape[1]) - - epsilon_in = torch.normal( - mean=torch.zeros([in_size]), std=torch.ones([in_size])) - epsilon_out = torch.normal( - mean=torch.zeros([out_size]), std=torch.ones([out_size])) - epsilon_in = self._f_epsilon(epsilon_in) - epsilon_out = self._f_epsilon(epsilon_out) - epsilon_w = torch.matmul( - torch.unsqueeze(epsilon_in, -1), - other=torch.unsqueeze(epsilon_out, 0)) - epsilon_b = epsilon_out - - sigma_w = torch.Tensor( - data=np.random.uniform( - low=-1.0 / np.sqrt(float(in_size)), - high=1.0 / np.sqrt(float(in_size)), - size=[in_size, out_size]), - dtype=torch.float32, - requires_grad=True) - # TF noise generation can be unreliable on GPU - # If generating the noise on the CPU, - # lowering sigma0 to 0.1 may be helpful - sigma_b = torch.Tensor( - data=np.full( - shape=[out_size], fill_value=sigma0 / np.sqrt(float(in_size))), - requires_grad=True) - w = torch.Tensor( - data=np.full( - shape=[in_size, out_size], - fill_value=6 / np.sqrt(float(in_size) + float(out_size))), - requires_grad=True) - b = torch.Tensor(data=np.zeros([out_size]), requires_grad=True) - action_activation = torch.matmul(action_in, w + sigma_w * epsilon_w) \ - + b + sigma_b * epsilon_b - - if not non_linear: - return action_activation - return nn.functional.relu(action_activation) - - def _f_epsilon(self, x): - return torch.sign(x) * torch.pow(torch.abs(x), 0.5) diff --git a/rllib/agents/dqn/dqn_torch_policy.py b/rllib/agents/dqn/dqn_torch_policy.py index 8ab1fe8bc..2965d1ac2 100644 --- a/rllib/agents/dqn/dqn_torch_policy.py +++ b/rllib/agents/dqn/dqn_torch_policy.py @@ -14,7 +14,8 @@ from ray.rllib.policy.torch_policy_template import build_torch_policy from ray.rllib.utils.error import UnsupportedSpaceException from ray.rllib.utils.exploration.parameter_noise import ParameterNoise from ray.rllib.utils.framework import try_import_torch -from ray.rllib.utils.torch_ops import huber_loss, reduce_mean_ignore_inf +from ray.rllib.utils.torch_ops import huber_loss, reduce_mean_ignore_inf, \ + softmax_cross_entropy_with_logits torch, nn = try_import_torch() F = None @@ -25,7 +26,9 @@ if nn: class QLoss: def __init__(self, q_t_selected, + q_logits_t_selected, q_tp1_best, + q_probs_tp1_best, importance_weights, rewards, done_mask, @@ -36,24 +39,61 @@ class QLoss: v_max=10.0): if num_atoms > 1: - raise ValueError("Torch version of DQN does not support " - "distributional Q yet!") + # Distributional Q-learning which corresponds to an entropy loss + z = torch.range(0.0, num_atoms - 1, dtype=torch.float32) + z = v_min + z * (v_max - v_min) / float(num_atoms - 1) - q_tp1_best_masked = (1.0 - done_mask) * q_tp1_best + # (batch_size, 1) * (1, num_atoms) = (batch_size, num_atoms) + r_tau = torch.unsqueeze( + rewards, -1) + gamma**n_step * torch.unsqueeze( + 1.0 - done_mask, -1) * torch.unsqueeze(z, 0) + r_tau = torch.clamp(r_tau, v_min, v_max) + b = (r_tau - v_min) / ((v_max - v_min) / float(num_atoms - 1)) + lb = torch.floor(b) + ub = torch.ceil(b) - # compute RHS of bellman equation - q_t_selected_target = rewards + gamma**n_step * q_tp1_best_masked + # Indispensable judgement which is missed in most implementations + # when b happens to be an integer, lb == ub, so pr_j(s', a*) will + # be discarded because (ub-b) == (b-lb) == 0. + floor_equal_ceil = (ub - lb < 0.5).float() - # compute the error (potentially clipped) - self.td_error = q_t_selected - q_t_selected_target.detach() - self.loss = torch.mean( - importance_weights.float() * huber_loss(self.td_error)) - self.stats = { - "mean_q": torch.mean(q_t_selected), - "min_q": torch.min(q_t_selected), - "max_q": torch.max(q_t_selected), - "mean_td_error": torch.mean(self.td_error), - } + # (batch_size, num_atoms, num_atoms) + l_project = F.one_hot(lb.long(), num_atoms) + # (batch_size, num_atoms, num_atoms) + u_project = F.one_hot(ub.long(), num_atoms) + ml_delta = q_probs_tp1_best * (ub - b + floor_equal_ceil) + mu_delta = q_probs_tp1_best * (b - lb) + ml_delta = torch.sum( + l_project * torch.unsqueeze(ml_delta, -1), dim=1) + mu_delta = torch.sum( + u_project * torch.unsqueeze(mu_delta, -1), dim=1) + m = ml_delta + mu_delta + + # Rainbow paper claims that using this cross entropy loss for + # priority is robust and insensitive to `prioritized_replay_alpha` + self.td_error = softmax_cross_entropy_with_logits( + logits=q_logits_t_selected, labels=m) + self.loss = torch.mean(self.td_error * importance_weights) + self.stats = { + # TODO: better Q stats for dist dqn + "mean_td_error": torch.mean(self.td_error), + } + else: + q_tp1_best_masked = (1.0 - done_mask) * q_tp1_best + + # compute RHS of bellman equation + q_t_selected_target = rewards + gamma**n_step * q_tp1_best_masked + + # compute the error (potentially clipped) + self.td_error = q_t_selected - q_t_selected_target.detach() + self.loss = torch.mean( + importance_weights.float() * huber_loss(self.td_error)) + self.stats = { + "mean_q": torch.mean(q_t_selected), + "min_q": torch.min(q_t_selected), + "max_q": torch.max(q_t_selected), + "mean_td_error": torch.mean(self.td_error), + } class ComputeTDErrorMixin: @@ -102,9 +142,12 @@ def build_q_model_and_distribution(policy, obs_space, action_space, config): framework="torch", model_interface=DQNTorchModel, name=Q_SCOPE, - dueling=config["dueling"], q_hiddens=config["hiddens"], + dueling=config["dueling"], + num_atoms=config["num_atoms"], use_noisy=config["noisy"], + v_min=config["v_min"], + v_max=config["v_max"], sigma0=config["sigma0"], # TODO(sven): Move option to add LayerNorm after each Dense # generically into ModelCatalog. @@ -120,9 +163,12 @@ def build_q_model_and_distribution(policy, obs_space, action_space, config): framework="torch", model_interface=DQNTorchModel, name=Q_TARGET_SCOPE, - dueling=config["dueling"], q_hiddens=config["hiddens"], + dueling=config["dueling"], + num_atoms=config["num_atoms"], use_noisy=config["noisy"], + v_min=config["v_min"], + v_max=config["v_max"], sigma0=config["sigma0"], # TODO(sven): Move option to add LayerNorm after each Dense # generically into ModelCatalog. @@ -149,54 +195,63 @@ def get_distribution_inputs_and_class(policy, def build_q_losses(policy, model, _, train_batch): config = policy.config - # q network evaluation - q_t = compute_q_values( + # Q-network evaluation. + q_t, q_logits_t, q_probs_t = compute_q_values( policy, policy.q_model, train_batch[SampleBatch.CUR_OBS], explore=False, is_training=True) - # target q network evalution - q_tp1 = compute_q_values( + # Target Q-network evaluation. + q_tp1, q_logits_tp1, q_probs_tp1 = compute_q_values( policy, policy.target_q_model, train_batch[SampleBatch.NEXT_OBS], explore=False, is_training=True) - # q scores for actions which we know were selected in the given state. - one_hot_selection = F.one_hot(train_batch[SampleBatch.ACTIONS], - policy.action_space.n) + # Q scores for actions which we know were selected in the given state. + one_hot_selection = F.one_hot( + train_batch[SampleBatch.ACTIONS], policy.action_space.n) q_t_selected = torch.sum( torch.where(q_t > -float("inf"), q_t, torch.tensor(0.0)) * one_hot_selection, 1) + q_logits_t_selected = torch.sum( + q_logits_t * torch.unsqueeze(one_hot_selection, -1), 1) # compute estimate of best possible value starting from state at t + 1 if config["double_q"]: - q_tp1_using_online_net = compute_q_values( - policy, - policy.q_model, - train_batch[SampleBatch.NEXT_OBS], - explore=False, - is_training=True) + q_tp1_using_online_net, q_logits_tp1_using_online_net, \ + q_dist_tp1_using_online_net = compute_q_values( + policy, + policy.q_model, + train_batch[SampleBatch.NEXT_OBS], + explore=False, + is_training=True) q_tp1_best_using_online_net = torch.argmax(q_tp1_using_online_net, 1) - q_tp1_best_one_hot_selection = F.one_hot(q_tp1_best_using_online_net, - policy.action_space.n) + q_tp1_best_one_hot_selection = F.one_hot( + q_tp1_best_using_online_net, policy.action_space.n) + q_tp1_best = torch.sum( + torch.where(q_tp1 > -float("inf"), q_tp1, torch.tensor(0.0)) * + q_tp1_best_one_hot_selection, 1) + q_probs_tp1_best = torch.sum( + q_probs_tp1 * torch.unsqueeze(q_tp1_best_one_hot_selection, -1), 1) else: q_tp1_best_one_hot_selection = F.one_hot( torch.argmax(q_tp1, 1), policy.action_space.n) + q_tp1_best = torch.sum( + torch.where(q_tp1 > -float("inf"), q_tp1, torch.tensor(0.0)) * + q_tp1_best_one_hot_selection, 1) + q_probs_tp1_best = torch.sum( + q_probs_tp1 * torch.unsqueeze(q_tp1_best_one_hot_selection, -1), 1) - q_tp1_best = torch.sum( - torch.where(q_tp1 > -float("inf"), q_tp1, torch.tensor(0.0)) * - q_tp1_best_one_hot_selection, 1) - - policy.q_loss = QLoss(q_t_selected, q_tp1_best, train_batch[PRIO_WEIGHTS], - train_batch[SampleBatch.REWARDS], - train_batch[SampleBatch.DONES].float(), - config["gamma"], config["n_step"], - config["num_atoms"], config["v_min"], - config["v_max"]) + policy.q_loss = QLoss( + q_t_selected, q_logits_t_selected, q_tp1_best, q_probs_tp1_best, + train_batch[PRIO_WEIGHTS], train_batch[SampleBatch.REWARDS], + train_batch[SampleBatch.DONES].float(), config["gamma"], + config["n_step"], config["num_atoms"], + config["v_min"], config["v_max"]) return policy.q_loss.loss @@ -225,26 +280,44 @@ def after_init(policy, obs_space, action_space, config): def compute_q_values(policy, model, obs, explore, is_training=False): - if policy.config["num_atoms"] > 1: - raise ValueError("torch DQN does not support distributional DQN yet!") + config = policy.config model_out, state = model({ SampleBatch.CUR_OBS: obs, "is_training": is_training, }, [], None) - advantages_or_q_values = model.get_advantages_or_q_values(model_out) - - if policy.config["dueling"]: - state_value = model.get_state_value(model_out) - advantages_mean = reduce_mean_ignore_inf(advantages_or_q_values, 1) - advantages_centered = advantages_or_q_values - torch.unsqueeze( - advantages_mean, 1) - q_values = state_value + advantages_centered + if config["num_atoms"] > 1: + (action_scores, z, support_logits_per_action, logits, + probs_or_logits) = model.get_q_value_distributions(model_out) else: - q_values = advantages_or_q_values + (action_scores, logits, + probs_or_logits) = model.get_q_value_distributions(model_out) - return q_values + if config["dueling"]: + state_score = model.get_state_value(model_out) + if policy.config["num_atoms"] > 1: + support_logits_per_action_mean = torch.mean( + support_logits_per_action, dim=1) + support_logits_per_action_centered = ( + support_logits_per_action - torch.unsqueeze( + support_logits_per_action_mean, dim=1)) + support_logits_per_action = torch.unsqueeze( + state_score, dim=1) + support_logits_per_action_centered + support_prob_per_action = nn.functional.softmax( + support_logits_per_action) + value = torch.sum(z * support_prob_per_action, dim=-1) + logits = support_logits_per_action + probs_or_logits = support_prob_per_action + else: + advantages_mean = reduce_mean_ignore_inf(action_scores, 1) + advantages_centered = action_scores - torch.unsqueeze( + advantages_mean, 1) + value = state_score + advantages_centered + else: + value = action_scores + + return value, logits, probs_or_logits def grad_process_and_td_error_fn(policy, optimizer, loss): diff --git a/rllib/agents/dqn/tests/test_dqn.py b/rllib/agents/dqn/tests/test_dqn.py index 87db3eaa8..5f8fe368a 100644 --- a/rllib/agents/dqn/tests/test_dqn.py +++ b/rllib/agents/dqn/tests/test_dqn.py @@ -22,8 +22,9 @@ class TestDQN(unittest.TestCase): config["num_workers"] = 2 num_iterations = 1 - for fw in framework_iterator(config): + for _ in framework_iterator(config): # Double-dueling DQN. + print("Double-dueling") plain_config = config.copy() trainer = dqn.DQNTrainer(config=plain_config, env="CartPole-v0") for i in range(num_iterations): @@ -34,9 +35,7 @@ class TestDQN(unittest.TestCase): trainer.stop() # Rainbow. - # TODO(sven): Add torch once DQN-torch supports distributional-Q. - if fw == "torch": - continue + print("Rainbow") rainbow_config = config.copy() rainbow_config["num_atoms"] = 10 rainbow_config["noisy"] = True diff --git a/rllib/examples/env/env_with_subprocess.py b/rllib/examples/env/env_with_subprocess.py index e55e8e0e7..6fccf3acd 100644 --- a/rllib/examples/env/env_with_subprocess.py +++ b/rllib/examples/env/env_with_subprocess.py @@ -3,14 +3,13 @@ import gym from gym.spaces import Discrete import os import subprocess -import time class EnvWithSubprocess(gym.Env): - """Our env that spawns a subprocess.""" + """An env that spawns a subprocess.""" # Dummy command to run as a subprocess with a unique name - UNIQUE_CMD = "sleep {}".format(str(time.time())) + UNIQUE_CMD = "sleep 20" def __init__(self, config): self.UNIQUE_FILE_0 = config["tmp_file1"] @@ -20,11 +19,11 @@ class EnvWithSubprocess(gym.Env): self.action_space = Discrete(2) self.observation_space = Discrete(2) - # Subprocess that should be cleaned up + # Subprocess that should be cleaned up. self.subproc = subprocess.Popen( self.UNIQUE_CMD.split(" "), shell=False) self.config = config - # Exit handler should be called + # Exit handler should be called. atexit.register(lambda: self.subproc.kill()) if config.worker_index == 0: atexit.register(lambda: os.unlink(self.UNIQUE_FILE_0)) diff --git a/rllib/models/torch/misc.py b/rllib/models/torch/misc.py index 27cf4c964..61a5ac972 100644 --- a/rllib/models/torch/misc.py +++ b/rllib/models/torch/misc.py @@ -1,7 +1,7 @@ """ Code adapted from https://github.com/ikostrikov/pytorch-a3c""" import numpy as np -from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.framework import get_activation_fn, try_import_torch torch, nn = try_import_torch() @@ -19,17 +19,21 @@ def same_padding(in_size, filter_size, stride_size): """Note: Padding is added to match TF conv2d `same` padding. See www.tensorflow.org/versions/r0.12/api_docs/python/nn/convolution - Params: + Args: in_size (tuple): Rows (Height), Column (Width) for input - stride_size (tuple): Rows (Height), Column (Width) for stride - filter_size (tuple): Rows (Height), Column (Width) for filter + stride_size (Union[int,Tuple[int, int]]): Rows (Height), column (Width) + for stride. If int, height == width. + filter_size (tuple): Rows (Height), column (Width) for filter - Output: + Returns: padding (tuple): For input into torch.nn.ZeroPad2d. output (tuple): Output shape after padding and convolution. """ in_height, in_width = in_size - filter_height, filter_width = filter_size + if isinstance(filter_size, int): + filter_height, filter_width = filter_size, filter_size + else: + filter_height, filter_width = filter_size stride_height, stride_width = stride_size out_height = np.ceil(float(in_height) / float(stride_height)) @@ -102,7 +106,9 @@ class SlimFC(nn.Module): if use_bias is True: nn.init.constant_(linear.bias, bias_init) layers.append(linear) - if activation_fn: + if isinstance(activation_fn, str): + activation_fn = get_activation_fn(activation_fn, "torch") + if activation_fn is not None: layers.append(activation_fn()) self._model = nn.Sequential(*layers) diff --git a/rllib/models/torch/modules/noisy_layer.py b/rllib/models/torch/modules/noisy_layer.py new file mode 100644 index 000000000..2e10d8e0d --- /dev/null +++ b/rllib/models/torch/modules/noisy_layer.py @@ -0,0 +1,90 @@ +import numpy as np + +from ray.rllib.utils.framework import get_activation_fn, try_import_torch +from ray.rllib.utils.framework import get_variable + +torch, nn = try_import_torch() + + +class NoisyLayer(nn.Module): + """A Layer that adds learnable Noise + a common dense layer: y = w^{T}x + b + a noisy layer: y = (w + \\epsilon_w*\\sigma_w)^{T}x + + (b+\\epsilon_b*\\sigma_b) + where \epsilon are random variables sampled from factorized normal + distributions and \\sigma are trainable variables which are expected to + vanish along the training procedure + """ + + def __init__(self, in_size, out_size, sigma0, activation="relu"): + """Initializes a NoisyLayer object. + + Args: + in_size: + out_size: + sigma0: + non_linear: + """ + super().__init__() + + self.in_size = in_size + self.out_size = out_size + self.sigma0 = sigma0 + self.activation = get_activation_fn(activation, framework="torch") + if self.activation is not None: + self.activation = self.activation() + + self.sigma_w = get_variable( + np.random.uniform( + low=-1.0 / np.sqrt(float(self.in_size)), + high=1.0 / np.sqrt(float(self.in_size)), + size=[self.in_size, out_size]), + framework="torch", + dtype=torch.float32, + torch_tensor=True, + trainable=True) + self.sigma_b = get_variable( + np.full( + shape=[out_size], + fill_value=sigma0 / np.sqrt(float(self.in_size))), + framework="torch", + dtype=torch.float32, + torch_tensor=True, + trainable=True) + self.w = get_variable( + np.full( + shape=[self.in_size, self.out_size], + fill_value=6 / np.sqrt(float(in_size) + float(out_size))), + framework="torch", + dtype=torch.float32, + torch_tensor=True, + trainable=True) + self.b = get_variable( + np.zeros([out_size]), + framework="torch", + dtype=torch.float32, + torch_tensor=True, + trainable=True) + + def forward(self, inputs): + epsilon_in = self._f_epsilon(torch.normal( + mean=torch.zeros([self.in_size]), + std=torch.ones([self.in_size]))) + epsilon_out = self._f_epsilon(torch.normal( + mean=torch.zeros([self.out_size]), + std=torch.ones([self.out_size]))) + epsilon_w = torch.matmul( + torch.unsqueeze(epsilon_in, -1), + other=torch.unsqueeze(epsilon_out, 0)) + epsilon_b = epsilon_out + + action_activation = torch.matmul( + inputs, self.w + self.sigma_w * epsilon_w + ) + self.b + self.sigma_b * epsilon_b + + if self.activation is not None: + action_activation = self.activation(action_activation) + return action_activation + + def _f_epsilon(self, x): + return torch.sign(x) * torch.pow(torch.abs(x), 0.5) diff --git a/rllib/tests/test_rollout_worker.py b/rllib/tests/test_rollout_worker.py index 6c56c66f7..ccd2280f3 100644 --- a/rllib/tests/test_rollout_worker.py +++ b/rllib/tests/test_rollout_worker.py @@ -267,7 +267,7 @@ class TestRolloutWorker(unittest.TestCase): pg.stop() def test_reward_clipping(self): - # clipping on + # Clipping: on. ev = RolloutWorker( env_creator=lambda _: MockEnv2(episode_length=10), policy=MockPolicy, @@ -278,7 +278,7 @@ class TestRolloutWorker(unittest.TestCase): self.assertEqual(result["episode_reward_mean"], 1000) ev.stop() - # clipping off + # Clipping: off. ev2 = RolloutWorker( env_creator=lambda _: MockEnv2(episode_length=10), policy=MockPolicy, diff --git a/rllib/utils/framework.py b/rllib/utils/framework.py index 6a140c1a7..5763e8169 100644 --- a/rllib/utils/framework.py +++ b/rllib/utils/framework.py @@ -206,6 +206,10 @@ def get_variable(value, elif framework == "torch" and torch_tensor is True: torch, _ = try_import_torch() var_ = torch.from_numpy(value) + if dtype == torch.float32: + var_ = var_.float() + elif dtype == torch.int32: + var_ = var_.int() if device: var_ = var_.to(device) var_.requires_grad = trainable diff --git a/rllib/utils/torch_ops.py b/rllib/utils/torch_ops.py index f4da0a5d6..42e1ad69b 100644 --- a/rllib/utils/torch_ops.py +++ b/rllib/utils/torch_ops.py @@ -4,7 +4,11 @@ import tree from ray.rllib.models.repeated_values import RepeatedValues from ray.rllib.utils.framework import try_import_torch -torch, _ = try_import_torch() +torch, nn = try_import_torch() + + +def atanh(x): + return 0.5 * torch.log((1 + x) / (1 - x)) def explained_variance(y, pred): @@ -50,13 +54,6 @@ def l2_loss(x): return torch.sum(torch.pow(x, 2.0)) / 2.0 -def reduce_mean_ignore_inf(x, axis): - """Same as torch.mean() but ignores -inf values.""" - mask = torch.ne(x, float("-inf")) - x_zeroed = torch.where(mask, x, torch.zeros_like(x)) - return torch.sum(x_zeroed, axis) / torch.sum(mask.float(), axis) - - def minimize_and_clip(optimizer, clip_val=10): """Clips gradients found in `optimizer.param_groups` to given value. @@ -69,6 +66,13 @@ def minimize_and_clip(optimizer, clip_val=10): torch.nn.utils.clip_grad_norm_(p.grad, clip_val) +def reduce_mean_ignore_inf(x, axis): + """Same as torch.mean() but ignores -inf values.""" + mask = torch.ne(x, float("-inf")) + x_zeroed = torch.where(mask, x, torch.zeros_like(x)) + return torch.sum(x_zeroed, axis) / torch.sum(mask.float(), axis) + + def sequence_mask(lengths, maxlen=None, dtype=None): """Offers same behavior as tf.sequence_mask for torch. @@ -86,6 +90,18 @@ def sequence_mask(lengths, maxlen=None, dtype=None): return mask +def softmax_cross_entropy_with_logits(logits, labels): + """Same behavior as tf.nn.softmax_cross_entropy_with_logits. + + Args: + x (TensorType): + + Returns: + + """ + return torch.sum(-labels * nn.functional.log_softmax(logits, -1), -1) + + def convert_to_non_torch_type(stats): """Converts values in `stats` to non-Tensor numpy or python types. @@ -138,7 +154,3 @@ def convert_to_torch_tensor(x, device=None): return tensor if device is None else tensor.to(device) return tree.map_structure(mapping, x) - - -def atanh(x): - return 0.5 * torch.log((1 + x) / (1 - x))