Remove state from act_func arguments

This commit is contained in:
Djordje Grbic
2024-04-16 12:18:22 +02:00
parent 5fdf33ea14
commit 4336c81ae3

View File

@@ -39,7 +39,7 @@ class BraxEnv(RLEnv):
def step(key, env_state, obs):
key, _ = jax.random.split(key)
action = act_func(state, obs, params)
action = act_func(obs, params)
obs, env_state, r, done, _ = self.step(randkey, env_state, action)
return key, env_state, obs, r, done