add action_policy for problem;

tensorneat/problem/rl_env/brax_env.py

@@ -5,8 +5,10 @@ from .rl_jit import RLEnv
 
 
 class BraxEnv(RLEnv):
-    def __init__(self, max_step=1000, repeat_times=1, record_episode=False, env_name: str = "ant", backend: str = "generalized"):
-        super().__init__(max_step, repeat_times, record_episode)
+    def __init__(
+        self, env_name: str = "ant", backend: str = "generalized", *args, **kwargs
+    ):
+        super().__init__(*args, **kwargs)
         self.env = envs.create(env_name=env_name, backend=backend)
 
     def env_step(self, randkey, env_state, action):

tensorneat/problem/rl_env/gymnax_env.py

@@ -4,8 +4,8 @@ from .rl_jit import RLEnv
 
 
 class GymNaxEnv(RLEnv):
-    def __init__(self, env_name, max_step=1000, repeat_times=1, record_episode=False):
-        super().__init__(max_step, repeat_times, record_episode)
+    def __init__(self, env_name, *args, **kwargs):
+        super().__init__(*args, **kwargs)
         assert env_name in gymnax.registered_envs, f"Env {env_name} not registered"
         self.env, self.env_params = gymnax.make(env_name)
 

tensorneat/problem/rl_env/jumanji/jumanji_2048.py

@@ -7,14 +7,21 @@ from ..rl_jit import RLEnv
 
 class Jumanji_2048(RLEnv):
     def __init__(
-        self, max_step=1000, repeat_times=1, record_episode=False, guarantee_invalid_action=True
+        self, guarantee_invalid_action=True, *args, **kwargs
     ):
-        super().__init__(max_step, repeat_times, record_episode)
+        super().__init__(*args, **kwargs)
         self.guarantee_invalid_action = guarantee_invalid_action
         self.env = jumanji.make("Game2048-v1")
 
     def env_step(self, randkey, env_state, action):
         action_mask = env_state["action_mask"]
+
+        ###################################################################
+
+        action = jnp.concatenate([action, jnp.full((4 - action.shape[0], ), -99999)])
+        action = (action - 1) / 15
+
+        ###################################################################
         if self.guarantee_invalid_action:
             score_with_mask = jnp.where(action_mask, action, -jnp.inf)
             action = jnp.argmax(score_with_mask)

tensorneat/problem/rl_env/rl_jit.py

@@ -11,11 +11,18 @@ from .. import BaseProblem
 class RLEnv(BaseProblem):
     jitable = True
 
-    def __init__(self, max_step=1000, repeat_times=1, record_episode=False):
+    def __init__(
+        self,
+        max_step=1000,
+        repeat_times=1,
+        record_episode=False,
+        action_policy: Callable = None,
+    ):
         super().__init__()
         self.max_step = max_step
         self.record_episode = record_episode
         self.repeat_times = repeat_times
+        self.action_policy = action_policy
 
     def evaluate(self, state: State, randkey, act_func: Callable, params):
         keys = jax.random.split(randkey, self.repeat_times)
@@ -63,7 +70,11 @@ class RLEnv(BaseProblem):
 
         def body_func(carry):
             obs, env_state, rng, done, tr, count, epis = carry  # tr -> total reward
-            action = act_func(state, params, obs)
+            if self.action_policy is not None:
+                forward_func = lambda obs: act_func(state, params, obs)
+                action = self.action_policy(forward_func, obs)
+            else:
+                action = act_func(state, params, obs)
             next_obs, next_env_state, reward, done, _ = self.step(
                 rng, env_state, action
             )