[RLlib] Allow for more than 2^31 policy timesteps. (#11301)

rllib/utils/exploration/epsilon_greedy.py

@@ -1,3 +1,4 @@
+import numpy as np
 from typing import Union, Optional
 
 from ray.rllib.models.action_dist import ActionDistribution
@@ -54,7 +55,10 @@ class EpsilonGreedy(Exploration):
 
         # The current timestep value (tf-var or python int).
         self.last_timestep = get_variable(
-            0, framework=framework, tf_name="timestep")
+            np.array(0, np.int64),
+            framework=framework,
+            tf_name="timestep",
+            dtype=np.int64)
 
         # Build the tf-info-op.
         if self.framework in ["tf2", "tf", "tfe"]:

rllib/utils/exploration/gaussian_noise.py

@@ -1,4 +1,5 @@
 from gym.spaces import Space
+import numpy as np
 from typing import Union, Optional
 
 from ray.rllib.models.action_dist import ActionDistribution
@@ -74,7 +75,10 @@ class GaussianNoise(Exploration):
 
         # The current timestep value (tf-var or python int).
         self.last_timestep = get_variable(
-            0, framework=self.framework, tf_name="timestep")
+            np.array(0, np.int64),
+            framework=self.framework,
+            tf_name="timestep",
+            dtype=np.int64)
 
         # Build the tf-info-op.
         if self.framework in ["tf2", "tf", "tfe"]:

rllib/utils/exploration/stochastic_sampling.py

@@ -1,4 +1,5 @@
 import gym
+import numpy as np
 import tree
 from typing import Union
 
@@ -53,7 +54,10 @@ class StochasticSampling(Exploration):
 
         # The current timestep value (tf-var or python int).
         self.last_timestep = get_variable(
-            0, framework=self.framework, tf_name="timestep")
+            np.array(0, np.int64),
+            framework=self.framework,
+            tf_name="timestep",
+            dtype=np.int64)
 
     @override(Exploration)
     def get_exploration_action(self,

rllib/utils/framework.py

@@ -1,4 +1,5 @@
 import logging
+import numpy as np
 import os
 import sys
 from typing import Any, Optional
@@ -185,6 +186,7 @@ def get_variable(value,
             does not have any (e.g. if it's an initializer w/o explicit value).
         dtype (Optional[TensorType]): An optional dtype to use iff `value` does
             not have any (e.g. if it's an initializer w/o explicit value).
+            This should always be a numpy dtype (e.g. np.float32, np.int64).
 
     Returns:
         any: A framework-specific variable (tf.Variable, torch.tensor, or
@@ -207,10 +209,13 @@ 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:
+        if dtype in [torch.float32, np.float32]:
             var_ = var_.float()
-        elif dtype == torch.int32:
+        elif dtype in [torch.int32, np.int32]:
             var_ = var_.int()
+        elif dtype in [torch.float64, np.float64]:
+            var_ = var_.double()
+
         if device:
             var_ = var_.to(device)
         var_.requires_grad = trainable

rllib/utils/schedules/piecewise_schedule.py

@@ -65,7 +65,7 @@ class PiecewiseSchedule(Schedule):
             "provided!"
 
         endpoints = tf.cast(
-            tf.stack([e[0] for e in self.endpoints] + [-1]), tf.int32)
+            tf.stack([e[0] for e in self.endpoints] + [-1]), tf.int64)
 
         # Create all possible interpolation results.
         results_list = []
@@ -79,6 +79,7 @@ class PiecewiseSchedule(Schedule):
 
         # Return correct results tensor depending on where we find t.
         def _cond(i, x):
+            x = tf.cast(x, tf.int64)
             return tf.logical_not(
                 tf.logical_or(
                     tf.equal(endpoints[i + 1], -1),
@@ -88,5 +89,5 @@ class PiecewiseSchedule(Schedule):
             return (i + 1, t)
 
         idx_and_t = tf.while_loop(_cond, _body,
-                                  [tf.constant(0, dtype=tf.int32), t])
+                                  [tf.constant(0, dtype=tf.int64), t])
         return results_list[idx_and_t[0]]