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

@@ -4,8 +4,8 @@ from .rl_jit import RLEnv
class GymNaxEnv(RLEnv):
def __init__(self, env_name, max_step=1000, record_episode=False):
super().__init__(max_step, record_episode)
def __init__(self, env_name, max_step=1000, repeat_times=1, record_episode=False):
super().__init__(max_step, repeat_times, record_episode)
assert env_name in gymnax.registered_envs, f"Env {env_name} not registered"
self.env, self.env_params = gymnax.make(env_name)