Fix issues related to new behavior of JAX DeviceArray.copy()

PiperOrigin-RevId: 438711926
google · Apr 1, 2022 · 1dd31c8 · 1dd31c8
1 parent 83252f9
commit 1dd31c8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/trax/rl/task.py b/trax/rl/task.py
@@ -271,6 +271,7 @@ def play(env, policy, dm_suite=False, max_steps=None, last_observation=None):
   cur_trajectory = Trajectory(last_observation)
   while not done and (max_steps is None or cur_step < max_steps):
     action, dist_inputs = policy(cur_trajectory)
+    action = np.asarray(action)
     step = env.step(action)
     if dm_suite:
       (observation, reward, done) = (