[RLLib] RE3 exploration algorithm TF2 framework support

rllib/utils/exploration/random_encoder.py

@@ -114,7 +114,7 @@ def compute_states_entropy(
     """
     obs_embeds_ = np.reshape(obs_embeds, [-1, embed_dim])
     dist = np.linalg.norm(obs_embeds_[:, None, :] - obs_embeds_[None, :, :], axis=-1)
-    return dist.argsort(axis=-1)[:, :k_nn][:, -1]
+    return dist.argsort(axis=-1)[:, :k_nn][:, -1].astype(np.float32)
 
 
 @PublicAPI
@@ -288,6 +288,6 @@ def _postprocess_tf(self, policy, sample_batch, tf_sess):
         else:
             obs_embeds = tf.stop_gradient(
                 self._encoder_net({SampleBatch.OBS: sample_batch[SampleBatch.OBS]})[0]
-            )
+            ).numpy()
         sample_batch[SampleBatch.OBS_EMBEDS] = obs_embeds
         return sample_batch

rllib/utils/exploration/tests/test_random_encoder.py

@@ -3,6 +3,7 @@
 
 import pytest
 import ray
+from ray.rllib.utils.test_utils import framework_iterator
 import ray.rllib.algorithms.ppo as ppo
 import ray.rllib.algorithms.sac as sac
 from ray.rllib.algorithms.callbacks import RE3UpdateCallbacks
@@ -48,17 +49,18 @@ class RE3Callbacks(RE3UpdateCallbacks, config["callbacks"]):
         }
 
         num_iterations = 30
-        algo = algo_cls(config=config)
-        learnt = False
-        for i in range(num_iterations):
-            result = algo.train()
-            print(result)
-            if result["episode_reward_max"] > -900.0:
-                print("Reached goal after {} iters!".format(i))
-                learnt = True
-                break
-        algo.stop()
-        self.assertTrue(learnt)
+        for _ in framework_iterator(config, frameworks=("tf", "tf2"), session=True):
+            algo = algo_cls(config=config)
+            learnt = False
+            for i in range(num_iterations):
+                result = algo.train()
+                print(result)
+                if result["episode_reward_max"] > -900.0:
+                    print("Reached goal after {} iters!".format(i))
+                    learnt = True
+                    break
+            algo.stop()
+            self.assertTrue(learnt)
 
     def test_re3_ppo(self):
         """Tests RE3 with PPO."""