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/brax_env.py
+++ b/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):