add jumanji env;

add repeat times for rl_env

tensorneat/problem/rl_env/brax_env.py

@@ -5,8 +5,8 @@ from .rl_jit import RLEnv
 
 
 class BraxEnv(RLEnv):
-    def __init__(self, max_step=1000, record_episode=False, env_name: str = "ant", backend: str = "generalized"):
-        super().__init__(max_step, record_episode)
+    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)
         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, record_episode=False):
-        super().__init__(max_step, record_episode)
+    def __init__(self, env_name, max_step=1000, repeat_times=1, record_episode=False):
+        super().__init__(max_step, repeat_times, record_episode)
         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

@@ -0,0 +1,56 @@
+import jax, jax.numpy as jnp
+import jumanji
+
+from utils import State
+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
+    ):
+        super().__init__(max_step, repeat_times, record_episode)
+        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"]
+        if self.guarantee_invalid_action:
+            score_with_mask = jnp.where(action_mask, action, -jnp.inf)
+            action = jnp.argmax(score_with_mask)
+        else:
+            action = jnp.argmax(action)
+
+        done = ~action_mask[action]
+
+        env_state, timestep = self.env.step(env_state, action)
+        reward = timestep["reward"]
+
+        board, action_mask = timestep["observation"]
+        extras = timestep["extras"]
+
+        done = done | (jnp.sum(action_mask) == 0)  # all actions of invalid
+
+        return board.reshape(-1), env_state, reward, done, extras
+
+    def env_reset(self, randkey):
+        env_state, timestep = self.env.reset(randkey)
+        step_type = timestep["step_type"]
+        reward = timestep["reward"]
+        discount = timestep["discount"]
+        observation = timestep["observation"]
+        extras = timestep["extras"]
+        board, action_mask = observation
+
+        return board.reshape(-1), env_state
+
+    @property
+    def input_shape(self):
+        return (16,)
+
+    @property
+    def output_shape(self):
+        return (4,)
+
+    def show(self, state, randkey, act_func, params, *args, **kwargs):
+        raise NotImplementedError("GymNax render must rely on gym 0.19.0(old version).")

tensorneat/problem/rl_env/rl_jit.py

@@ -1,20 +1,47 @@
 from functools import partial
+from typing import Callable
 
 import jax
 import jax.numpy as jnp
 
+from utils import State
 from .. import BaseProblem
 
 
 class RLEnv(BaseProblem):
     jitable = True
 
-    def __init__(self, max_step=1000, record_episode=False):
+    def __init__(self, max_step=1000, repeat_times=1, record_episode=False):
         super().__init__()
         self.max_step = max_step
         self.record_episode = record_episode
+        self.repeat_times = repeat_times
 
-    def evaluate(self, state, randkey, act_func, params):
+    def evaluate(self, state: State, randkey, act_func: Callable, params):
+        keys = jax.random.split(randkey, self.repeat_times)
+        if self.record_episode:
+            rewards, episodes = jax.vmap(
+                self.evaluate_once, in_axes=(None, 0, None, None)
+            )(state, keys, act_func, params)
+            episodes["obs"] = episodes["obs"].reshape(
+                self.max_step * self.repeat_times, *self.input_shape
+            )
+            episodes["action"] = episodes["action"].reshape(
+                self.max_step * self.repeat_times, *self.output_shape
+            )
+            episodes["reward"] = episodes["reward"].reshape(
+                self.max_step * self.repeat_times,
+            )
+
+            return rewards.mean(), episodes
+
+        else:
+            rewards = jax.vmap(self.evaluate_once, in_axes=(None, 0, None, None))(
+                state, keys, act_func, params
+            )
+            return rewards.mean()
+
+    def evaluate_once(self, state, randkey, act_func, params):
         rng_reset, rng_episode = jax.random.split(randkey)
         init_obs, init_env_state = self.reset(rng_reset)