how to set a trigger for when to record video open ai gym

Question

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.

Answer 1

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)