Remove state from act_func arguments
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user