uint64 discrete actions

Author: jakegrigsby

amago/agent.py

@@ -763,7 +763,7 @@ def get_actions(
                 actions = action_dists.mean
         # get intended gamma distribution (always in -1 idx)
         actions = actions[..., -1, :]
-        dtype = torch.uint8 if (self.discrete or self.multibinary) else torch.float32
+        dtype = torch.int64 if (self.discrete or self.multibinary) else torch.float32
         return actions.to(dtype=dtype), hidden_state
 
     @torch.no_grad()

amago/envs/env_utils.py

@@ -141,7 +141,7 @@ def action(self, action):
         if isinstance(action, int):
             return action
         elif isinstance(action, np.ndarray) and action.ndim == 1:
-            assert action.shape[0] == 1 and action.dtype == np.uint8
+            assert action.shape[0] == 1
             return action[0].item()
         elif isinstance(action, np.ndarray) and action.ndim == 2:
             assert action.shape[1] == 1

amago/envs/exploration.py

@@ -205,7 +205,7 @@ def add_exploration_noise(self, action: np.ndarray, local_step: np.ndarray):
             use_random = np.expand_dims(self.rng.random(self.batched_envs) <= noise, 1)
             expl_action = (
                 use_random * random_action + (1 - use_random) * action
-            ).astype(np.uint8)
+            ).astype(np.int64)
             assert expl_action.shape == (self.batched_envs, 1)
         else:
             # random noise (TD3-style)