add args record_episode in rl tasks, with related test "test_record_episode.ipynb";

add args return_data in func_fit tasks.

tensorneat/problem/func_fit/func_fit.py

@@ -8,11 +8,12 @@ from .. import BaseProblem
 class FuncFit(BaseProblem):
     jitable = True
 
-    def __init__(self, error_method: str = "mse"):
+    def __init__(self, error_method: str = "mse", return_data: bool = False):
         super().__init__()
 
         assert error_method in {"mse", "rmse", "mae", "mape"}
         self.error_method = error_method
+        self.return_data = return_data
 
     def setup(self, state: State = State()):
         return state
@@ -38,7 +39,10 @@ class FuncFit(BaseProblem):
         else:
             raise NotImplementedError
 
-        return -loss
+        if self.return_data:
+            return -loss, self.inputs
+        else:
+            return -loss
 
     def show(self, state, randkey, act_func, params, *args, **kwargs):
         predict = jax.vmap(act_func, in_axes=(None, 0, None))(

tensorneat/problem/func_fit/xor.py

@@ -4,8 +4,6 @@ from .func_fit import FuncFit
 
 
 class XOR(FuncFit):
-    def __init__(self, error_method: str = "mse"):
-        super().__init__(error_method)
 
     @property
     def inputs(self):

tensorneat/problem/func_fit/xor3d.py

@@ -4,9 +4,6 @@ from .func_fit import FuncFit
 
 
 class XOR3d(FuncFit):
-    def __init__(self, error_method: str = "mse"):
-        super().__init__(error_method)
-
     @property
     def inputs(self):
         return np.array(

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, env_name: str = "ant", backend: str = "generalized"):
-        super().__init__(max_step)
+    def __init__(self, max_step=1000, record_episode=False, env_name: str = "ant", backend: str = "generalized"):
+        super().__init__(max_step, 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):
-        super().__init__(max_step)
+    def __init__(self, env_name, max_step=1000, record_episode=False):
+        super().__init__(max_step, 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/rl_jit.py

@@ -1,6 +1,7 @@
 from functools import partial
 
 import jax
+import jax.numpy as jnp
 
 from .. import BaseProblem
 
@@ -8,32 +9,64 @@ from .. import BaseProblem
 class RLEnv(BaseProblem):
     jitable = True
 
-    def __init__(self, max_step=1000):
+    def __init__(self, max_step=1000, record_episode=False):
         super().__init__()
         self.max_step = max_step
+        self.record_episode = record_episode
 
     def evaluate(self, state, randkey, act_func, params):
         rng_reset, rng_episode = jax.random.split(randkey)
         init_obs, init_env_state = self.reset(rng_reset)
 
+        if self.record_episode:
+            obs_array = jnp.full((self.max_step, *self.input_shape), jnp.nan)
+            action_array = jnp.full((self.max_step, *self.output_shape), jnp.nan)
+            reward_array = jnp.full((self.max_step,), jnp.nan)
+            episode = {
+                "obs": obs_array,
+                "action": action_array,
+                "reward": reward_array,
+            }
+        else:
+            episode = None
+
         def cond_func(carry):
-            _, _, _, done, _, count = carry
+            _, _, _, done, _, count, _ = carry
             return ~done & (count < self.max_step)
 
         def body_func(carry):
-            obs, env_state, rng, done, tr, count = carry  # tr -> total reward
+            obs, env_state, rng, done, tr, count, epis = carry  # tr -> total reward
             action = act_func(state, obs, params)
             next_obs, next_env_state, reward, done, _ = self.step(
                 rng, env_state, action
             )
             next_rng, _ = jax.random.split(rng)
-            return next_obs, next_env_state, next_rng, done, tr + reward, count + 1
 
-        _, _, _, _, total_reward, _ = jax.lax.while_loop(
-            cond_func, body_func, (init_obs, init_env_state, rng_episode, False, 0.0, 0)
+            if self.record_episode:
+                epis["obs"] = epis["obs"].at[count].set(obs)
+                epis["action"] = epis["action"].at[count].set(action)
+                epis["reward"] = epis["reward"].at[count].set(reward)
+
+            return (
+                next_obs,
+                next_env_state,
+                next_rng,
+                done,
+                tr + reward,
+                count + 1,
+                epis,
+            )
+
+        _, _, _, _, total_reward, _, episode = jax.lax.while_loop(
+            cond_func,
+            body_func,
+            (init_obs, init_env_state, rng_episode, False, 0.0, 0, episode),
         )
 
-        return total_reward
+        if self.record_episode:
+            return total_reward, episode
+        else:
+            return total_reward
 
     # @partial(jax.jit, static_argnums=(0,))
     def step(self, randkey, env_state, action):