[RLlib] Fix issue 8135 (DDPG inf actions when using [-inf,inf] action space). (#8302)

rllib/agents/ddpg/ddpg_tf_model.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from ray.rllib.models.tf.tf_modelv2 import TFModelV2
 from ray.rllib.utils import try_import_tf
 
@@ -53,6 +55,8 @@ class DDPGTFModel(TFModelV2):
 
         self.model_out = tf.keras.layers.Input(
             shape=(num_outputs, ), name="model_out")
+        self.bounded = np.logical_and(action_space.bounded_above,
+                                      action_space.bounded_below).any()
         self.action_dim = action_space.shape[0]
 
         if actor_hiddens:
@@ -72,17 +76,17 @@ class DDPGTFModel(TFModelV2):
 
         # Use sigmoid to scale to [0,1], but also double magnitude of input to
         # emulate behaviour of tanh activation used in DDPG and TD3 papers.
+        # After sigmoid squashing, re-scale to env action space bounds.
         def lambda_(x):
-            sigmoid_out = tf.nn.sigmoid(2 * x)
-            # Rescale to actual env policy scale
-            # (shape of sigmoid_out is [batch_size, dim_actions], so we reshape
-            # to get same dims)
             action_range = (action_space.high - action_space.low)[None]
             low_action = action_space.low[None]
-            actions = action_range * sigmoid_out + low_action
-            return actions
+            sigmoid_out = tf.nn.sigmoid(2 * x)
+            squashed = action_range * sigmoid_out + low_action
+            return squashed
 
-        actor_out = tf.keras.layers.Lambda(lambda_)(actor_out)
+        # Only squash if we have bounded actions.
+        if self.bounded:
+            actor_out = tf.keras.layers.Lambda(lambda_)(actor_out)
 
         self.policy_model = tf.keras.Model(self.model_out, actor_out)
         self.register_variables(self.policy_model.variables)

rllib/agents/ddpg/ddpg_torch_model.py

@@ -49,6 +49,11 @@ class DDPGTorchModel(TorchModelV2, nn.Module):
                               model_config, name)
         nn.Module.__init__(self)
 
+        self.bounded = np.logical_and(action_space.bounded_above,
+                                      action_space.bounded_below).any()
+        self.action_range = torch.from_numpy(
+            (action_space.high - action_space.low)[None])
+        self.low_action = torch.from_numpy(action_space.low[None])
         self.action_dim = np.product(action_space.shape)
 
         # Build the policy network.
@@ -81,19 +86,16 @@ class DDPGTorchModel(TorchModelV2, nn.Module):
 
         # Use sigmoid to scale to [0,1], but also double magnitude of input to
         # emulate behaviour of tanh activation used in DDPG and TD3 papers.
+        # After sigmoid squashing, re-scale to env action space bounds.
         class _Lambda(nn.Module):
-            def forward(self, x):
+            def forward(self_, x):
                 sigmoid_out = nn.Sigmoid()(2.0 * x)
-                # Rescale to actual env policy scale
-                # (shape of sigmoid_out is [batch_size, dim_actions],
-                # so we reshape to get same dims)
-                action_range = (action_space.high - action_space.low)[None]
-                low_action = action_space.low[None]
-                actions = torch.from_numpy(action_range) * sigmoid_out + \
-                    torch.from_numpy(low_action)
-                return actions
+                squashed = self.action_range * sigmoid_out + self.low_action
+                return squashed
 
-        self.policy_model.add_module("action_out_squashed", _Lambda())
+        # Only squash if we have bounded actions.
+        if self.bounded:
+            self.policy_model.add_module("action_out_squashed", _Lambda())
 
         # Build the Q-net(s), including target Q-net(s).
         def build_q_net(name_):

rllib/agents/ddpg/ddpg_torch_policy.py

@@ -80,6 +80,8 @@ def ddpg_actor_critic_loss(policy, model, _, train_batch):
     # Q-values for current policy (no noise) in given current state
     q_t_det_policy = model.get_q_values(model_out_t, policy_t)
 
+    actor_loss = -torch.mean(q_t_det_policy)
+
     if twin_q:
         twin_q_t = model.get_twin_q_values(model_out_t,
                                            train_batch[SampleBatch.ACTIONS])
@@ -127,7 +129,6 @@ def ddpg_actor_critic_loss(policy, model, _, train_batch):
             errors = 0.5 * torch.pow(td_error, 2.0)
 
     critic_loss = torch.mean(train_batch[PRIO_WEIGHTS] * errors)
-    actor_loss = -torch.mean(q_t_det_policy)
 
     # Add l2-regularization if required.
     if l2_reg is not None:
@@ -154,20 +155,23 @@ def ddpg_actor_critic_loss(policy, model, _, train_batch):
     policy.td_error = td_error
     policy.q_t = q_t
 
-    # Return one loss value (even though we treat them separately in our
-    # 2 optimizers: actor and critic).
+    # Return two loss terms (corresponding to the two optimizers, we create).
     return policy.actor_loss, policy.critic_loss
 
 
 def make_ddpg_optimizers(policy, config):
-    # Create separate optimizers for actor & critic losses.
+    """Create separate optimizers for actor & critic losses."""
+
+    # Set epsilons to match tf.keras.optimizers.Adam's epsilon default.
     policy._actor_optimizer = torch.optim.Adam(
         params=policy.model.policy_variables(),
         lr=config["actor_lr"],
-        eps=1e-7)  # to match tf.keras.optimizers.Adam's epsilon default
+        eps=1e-7)
+
     policy._critic_optimizer = torch.optim.Adam(
-        params=policy.model.q_variables(), lr=config["critic_lr"],
-        eps=1e-7)  # to match tf.keras.optimizers.Adam's epsilon default
+        params=policy.model.q_variables(), lr=config["critic_lr"], eps=1e-7)
+
+    # Return them in the same order as the respective loss terms are returned.
     return policy._actor_optimizer, policy._critic_optimizer
 
 

rllib/policy/torch_policy.py

@@ -232,6 +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)
 
         # Loop through all optimizers.
         grad_info = {"allreduce_latency": 0.0}
@@ -240,7 +241,6 @@ class TorchPolicy(Policy):
             opt.zero_grad()
             # Recompute gradients of loss over all variables.
             loss_out[i].backward(retain_graph=(i < len(self._optimizers) - 1))
-
             grad_info.update(self.extra_grad_process(opt, loss_out[i]))
 
             if self.distributed_world_size:

rllib/utils/exploration/ornstein_uhlenbeck_noise.py

@@ -96,8 +96,10 @@ class OrnsteinUhlenbeckNoise(GaussianNoise):
         ou_new = self.ou_theta * -self.ou_state + \
             self.ou_sigma * gaussian_sample
         ou_state_new = tf.assign_add(self.ou_state, ou_new)
-        noise = scale * self.ou_base_scale * ou_state_new * \
-            (self.action_space.high - self.action_space.low)
+        high_m_low = self.action_space.high - self.action_space.low
+        high_m_low = tf.where(
+            tf.math.is_inf(high_m_low), tf.ones_like(high_m_low), high_m_low)
+        noise = scale * self.ou_base_scale * ou_state_new * high_m_low
         stochastic_actions = tf.clip_by_value(
             deterministic_actions + noise,
             self.action_space.low * tf.ones_like(deterministic_actions),
@@ -156,6 +158,9 @@ class OrnsteinUhlenbeckNoise(GaussianNoise):
                 high_m_low = torch.from_numpy(
                     self.action_space.high - self.action_space.low). \
                     to(self.device)
+                high_m_low = torch.where(
+                    torch.isinf(high_m_low),
+                    torch.ones_like(high_m_low).to(self.device), high_m_low)
                 noise = scale * self.ou_base_scale * self.ou_state * high_m_low
                 action = torch.clamp(det_actions + noise,
                                      self.action_space.low[0],