Farama-Foundation / SuperSuit

A collection of wrappers for Gymnasium and PettingZoo environments (being merged into gymnasium.wrappers and pettingzoo.wrappers
Other
441 stars 56 forks source link

Render mode attribute is erased by ss wrappers #241

Closed Butanium closed 6 months ago

Butanium commented 7 months ago

This makes it impossible to use the sb3 video recorder on the fly for example:

from pettingzoo.atari import pong_v3
import supersuit as ss
from stable_baselines3.common.vec_env import VecVideoRecorder
def get_env():
    env = pong_v3.parallel_env(render_mode="rgb_array")
    env = ss.pettingzoo_env_to_vec_env_v1(env)
    envs = ss.concat_vec_envs_v1(
        env, 1 , num_cpus=0, base_class="stable_baselines3"
    )
    envs.render_mode = "rgb_array"
    envs = VecVideoRecorder(envs, f"videos/", capped_cubic_video_schedule)
    envs.single_observation_space = envs.observation_space
    envs.single_action_space = envs.action_space
    return envs

Without the envs.render_mode = "rgb_array" line, I get this error:

[<ipython-input-72-90f28677e07d>](https://localhost:8080/#) in get_env()
      6     )
      7     # envs.render_mode = "rgb_array"
----> 8     envs = VecVideoRecorder(envs, f"videos/", capped_cubic_video_schedule)
      9     envs.single_observation_space = envs.observation_space
     10     envs.single_action_space = envs.action_space

[/usr/local/lib/python3.10/dist-packages/stable_baselines3/common/vec_env/vec_video_recorder.py](https://localhost:8080/#) in __init__(self, venv, video_folder, record_video_trigger, video_length, name_prefix)
     50 
     51         self.env.metadata = metadata
---> 52         assert self.env.render_mode == "rgb_array", f"The render_mode must be 'rgb_array', not {self.env.render_mode}"
Butanium commented 6 months ago

I used

!pip install git+https://github.com/Farama-Foundation/PettingZoo/
!pip install git+https://github.com/Farama-Foundation/SuperSuit/

In this colab if you want to reproduce: https://colab.research.google.com/drive/14YWTsEtQk4oluCf-W9RXj077LZIyrkn4?usp=sharing

elliottower commented 6 months ago

Fixed in #243