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

add args record_episode in rl tasks, with related test "test_record_episode.ipynb";

Browse Source 
add args return_data in func_fit tasks.
This commit is contained in:
wls2002
2024-05-30 17:05:56 +08:00
parent 20320105e6
commit cd92f411dc
8 changed files with 512 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -5,8 +5,8 @@ from .rl_jit import RLEnv
class BraxEnv(RLEnv):
def __init__(self, max_step=1000, env_name: str = "ant", backend: str = "generalized"):
super().__init__(max_step)
def __init__(self, max_step=1000, record_episode=False, env_name: str = "ant", backend: str = "generalized"):
super().__init__(max_step, record_episode)
self.env = envs.create(env_name=env_name, backend=backend)
def env_step(self, randkey, env_state, action):