[Question] Render vector env

[sparisi]

Question

Say I have 10 environments and I'd like to render them in a grid-like manner (and eventually save a video like that).
How do I do it?
As I was running some pytest for a PR, I noticed that some tests were rendering environments like that, but I cannot see where/how? I checked tests file and docs about wrapper, but couldn't find anything.

The doc for RecordVideo says it doesn't support vectorized environments, so I guess that, even after rendering them in a grid-like manner, it is not possible to save a video of them?

[pseudo-rnd-thoughts]

The code is in the gymnasium.wrappres.vector.HumanRendering wrapper

It would be interesting to apply it to vector environments but with differing episode lengths etc, I'm not sure, what solution we would want

[sparisi]

The code is in the gymnasium.wrappres.vector.HumanRendering wrapper

It would be interesting to apply it to vector environments but with differing episode lengths etc, I'm not sure, what solution we would want

I see, so it's for rendering but not video recording right?

I made this utilities that I use for recording videos from VectorEnv, maybe they can be useful?
The way frames are concatenated into a grid seems different from the existing vector.HumanRendering. For example, when I ran pytest I noticed that 3 environments were rendered into a 2x2 grid with 1 black part. My code would render them into a 1x3 grid.

Let me know if this can be useful and (if so) where it could be added.
I have also attached a video to show the final video.
EDIT: GitHub video player seems broken for some reason, but you can see the video if you download it.

import numpy as np

"""
Concatenates a list of frames in a grid-like manner.
Useful to concatenate many videos and play them side-to-side.
Frames are concatenated taking their size into account, such that the ratio of
the concatenated frame is as close as possible to 1.

Example:

>>> import gymnasium
>>> import cv2
>>>
>>> n_envs = 20
>>> envs = [
>>>     gymnasium.make("Pendulum-v1", render_mode="rgb_array")
>>>     for _ in range(n_envs)
>>> ]
>>>
>>> for e in envs:
>>>     _ = e.reset()
>>>
>>> frames = []
>>> for t in range(100):
>>>     step_frames = []
>>>     for e in envs:
>>>         _ = e.step(e.action_space.sample())
>>>         step_frames.append(e.render())
>>>     frames.append(frames_to_grid(step_frames))
>>>
>>> out = cv2.VideoWriter(
>>>     "concat_videos.mp4",
>>>     cv2.VideoWriter_fourcc(*'mp4v'),
>>>     30,
>>>     (frames[0].shape[1], frames[0].shape[0]),
>>> )
>>>
>>> for frame in frames:
>>>     out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
>>>
>>> out.release()
"""

def reshape_frames_into_grid(n_frames, h, w, target_aspect_ratio=1.0):
    best_rows, best_cols = 1, n_frames
    min_diff = np.inf
    for rows in range(1, int(n_frames**0.5) + 1):
        if n_frames % rows == 0:
            cols = n_frames // rows
            total_height = rows * h
            total_width = cols * w
            aspect = total_width / total_height
            diff = abs(aspect - target_aspect_ratio)
            if diff < min_diff:
                min_diff = diff
                best_rows, best_cols = rows, cols
    return best_rows, best_cols

def frames_to_grid(frames, target_aspect_ratio=1.0):
    frames = np.array(frames)
    n_frames, h, w, c = frames.shape
    rows, cols = reshape_frames_into_grid(n_frames, h, w, target_aspect_ratio)
    grid = np.zeros((rows * h, cols * w, c), dtype=frames.dtype)
    for idx in range(n_frames):
        r = idx // cols
        c_ = idx % cols
        grid[r * h : (r + 1) * h, c_ * w : (c_ + 1) * w] = frames[idx]
    return grid

10000.mp4

[pseudo-rnd-thoughts]

@sparisi Thanks for that.

I think a vector RecordVideo wrapper would be really helpful.
Examining projects like CleanRL, I believe the capability to record only one of the sub-environments is crucial, along with the ability to record all or most environments.

Would you be able to build that out into a wrapper?
For the HumanRendering, I allowed users to specify a video shape, then we try to fit as many renders into the shape, which is why for 3 environments, there is a black spot as 2x2 is the most optimal fitting for the parameters but with recorded videos, we might want to allow different recording shapes if that makes sense

[sparisi]

@sparisi Thanks for that.

I think a vector RecordVideo wrapper would be really helpful. Examining projects like CleanRL, I believe the capability to record only one of the sub-environments is crucial, along with the ability to record all or most environments.

Would you be able to build that out into a wrapper? For the HumanRendering, I allowed users to specify a video shape, then we try to fit as many renders into the shape, which is why for 3 environments, there is a black spot as 2x2 is the most optimal fitting for the parameters but with recorded videos, we might want to allow different recording shapes if that makes sense

Sure, I'll work on that in my spare time. The best way is to create a new wrapper -- a vectorized version of RecordVideo, right?

[pseudo-rnd-thoughts]

Yes, a vectorized version of RecordVideo. Something simple to start with and we can get more complex

[sparisi]

@pseudo-rnd-thoughts I am finalizing the wrapper and writing some tests. What's the best way to handle resets in vectorized envs?

What I usually do is something like

ep_done = np.zeros((n_envs))
while not np.all(ep_done):
    act = env.action_space.sample()
    obs, rwd, term, trunc, _ = env.step(act)
    done = np.logical_or(term, trunc)
    ep_done = np.logical_or(done, ep_done)

This would be one episode.
Is there a cleaner way to reset when all are done?

[pseudo-rnd-thoughts]

Err, this is a very complex problem
Have you seen the blog post - https://farama.org/Vector-Autoreset-Mode

For all sub-environments, I see two problems

1. Sub-environments will almost always be at different stages of the episode, unless you reset them and record from there.
2. You need to choose how long to record for.

I would suggest that you choose either a particular sub-environment's episode or observation and record all sub-environments until the sub-environment ends or a frame limit is reached

Thoughts?

[sparisi]

Err, this is a very complex problem Have you seen the blog post - https://farama.org/Vector-Autoreset-Mode

For all sub-environments, I see two problems

1. Sub-environments will almost always be at different stages of the episode, unless you reset them and record from there.
2. You need to choose how long to record for.

I would suggest that you choose either a particular sub-environment's episode or observation and record all sub-environments until the sub-environment ends or a frame limit is reached

Thoughts?

Yes, I read it some time ago for the same reason (recording video).
For training, I think there is no practical solution to this problem, and one has to just manually check termination (like in the loop I posted).

For recording video, it would be great to have another mode where the env "stays still" if terminated or truncated, i.e., it just returns the same tuple (obs, rwd, term, trunc). This way, the frames in the video do not change and it's clear if an episode ended.
I don't think it's a big deal (having videos where the episode resets is fine, at least for the purpose of checking how the policy behaves). I just wanted to be sure to write tests properly (I'll use a step limit).