add jumanji env;

add repeat times for rl_env
This commit is contained in:
wls2002
2024-06-05 14:24:17 +08:00
parent edfb0596e7
commit 10ec1c2df9
10 changed files with 1615 additions and 7 deletions

View File

@@ -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):