add action_policy for problem;

2024-06-07 17:09:16 +08:00
parent 10ec1c2df9
commit 3d5b80c6fa
13 changed files with 2417 additions and 1191 deletions
--- a/tensorneat/problem/rl_env/gymnax_env.py
+++ b/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)