stable-baselines3: [Bug]: Can't create the PPO model on Macbook M1

🐛 Bug

I’m using Macbook Pro M1 and running code in Jupyter Notebook. Every time I run the scripts, I get an error and the Kernel crashed. I’ve tried running the code in Google Colab and it went well. It seems to me that stable-baselines3 library doesn’t support M1 chips. What do I need to do now to run the code locally?

To Reproduce

from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env

# Parallel environments
env = make_vec_env("CartPole-v1", n_envs=4)

model = PPO("MlpPolicy", env, verbose=1)
model.learn(total_timesteps=25000)
model.save("ppo_cartpole")

del model # remove to demonstrate saving and loading

model = PPO.load("ppo_cartpole")

obs = env.reset()
while True:
    action, _states = model.predict(obs)
    obs, rewards, dones, info = env.step(action)
    env.render()

Relevant log output / Error message

The Kernel crashed while executing code in the the current cell or a previous cell. Please review the code in the cell(s) to identify a possible cause of the failure. Click here for more info. View Jupyter log for further details.

Canceled future for execute_request message before replies were done

System Info

No response

Checklist

• I have checked that there is no similar issue in the repo
• I have read the documentation
• I have provided a minimal working example to reproduce the bug
• I’ve used the markdown code blocks for both code and stack traces.