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

fix show error of cartpolev1

Browse Source
This commit is contained in:
Zenbook
2025-02-17 23:34:33 +08:00
parent d205d34ad2
commit d86a3196bd
1 changed files with 10 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
src/tensorneat/problem/rl/brax.py
View File

@@ -41,7 +41,7 @@ class BraxEnv(RLEnv):
*args, *args,
**kwargs, **kwargs,
): ):
assert output_type in ["rgb_array", "gif"] assert output_type in ["rgb_array", "gif"]
import jax import jax
@@ -76,12 +76,17 @@ class BraxEnv(RLEnv):
reward += r reward += r
if done: if done:
break break
print("Total reward: ", reward) print("Total reward: ", reward)
imgs = image.render_array( try:
sys=self.env.sys, trajectory=state_histories, height=height, width=width, camera="track" imgs = image.render_array(
) sys=self.env.sys, trajectory=state_histories, height=height, width=width, camera="track"
)
except ValueError:
imgs = image.render_array(
sys=self.env.sys, trajectory=state_histories, height=height, width=width
)
if output_type == "rgb_array": if output_type == "rgb_array":
imgs = np.array(imgs) imgs = np.array(imgs)