I run multiple episodes, but only want to record for specific ones. More specifically, I'd like to have an input acting as trigger. Currently, my code is this:
env = gym.make("HalfCheetah-v3")
env = gym.wrappers.RecordVideo(env, 'video')
env.reset()
for t in range(200):
action = env.action_space.sample()
observation, reward, done, info = env.step(action)
env.close()
env.reset()
#I only one to record this run
for t in range(200):
action = env.action_space.sample()
observation, reward, done, info = env.step(action)
env.close()
In that case, there are two episodes, but I'd like to only record the second one. Or in general, I'd like to have a input where I can select which episodes are recorded and which episodes are not recorded.
The function RecordVideo has an argument which is step_trigger, which is of the form: step_trigger: Callable[[int], bool] = None
, but I don't really know how to use it.
I found this other question here, and they use lambda id_episode:True
. I tried that, but didn't make a difference because I probably wasn't using it correctly. I'd appreciate any help.
From the documentation of Wrappers on gym's website, the episode/ step trigger should be a function that accepts episode/ step index and returns a bool
value. Therefore, for example, if you want to record a video of the second episode only, the wrapper should be used like this:
#record video for the second episode
env = gym.wrappers.RecordVideo(env, 'video', episode_trigger = lambda x: x == 2)
or, if you want to record video for every even steps:
#record video for every even episodes
env = gym.wrappers.RecordVideo(env, 'video', episode_trigger = lambda x: x % 2 == 0)