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

add args return_data in func_fit tasks.

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)