tony/tensorneat-mend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Remove state from act_func arguments

Browse Source
This commit is contained in:
Djordje Grbic
2024-04-16 12:18:22 +02:00
parent 5fdf33ea14
commit 4336c81ae3
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tensorneat/problem/rl_env/brax_env.py
View File

@@ -39,7 +39,7 @@ class BraxEnv(RLEnv):
def step(key, env_state, obs): def step(key, env_state, obs):
key, _ = jax.random.split(key) 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) obs, env_state, r, done, _ = self.step(randkey, env_state, action)
return key, env_state, obs, r, done return key, env_state, obs, r, done