The right way to install stable-baselines?

I am trying to install stable-baselines and run the first two lines from Getting Started section of the online manual but no option is working. I started with

pip install stable-baselines

Now when I run:

import gym
from stable_baselines.common.policies import MlpPolicy

I get

No module named 'tensorflow.contrib'

This apparently is because tensorflow version 2 doesn't have tensorflow.contrib. But version 2 was released in Sept 2019. Do I really have to use only tensorflow version 1?

What is the right way to install stable-baselines and run that simple example?


I tried

pip install stable-baselines3 

in a virtual environment. This gives a different error:

In [2]: from stable_baselines.common.policies import MlpPolicy
---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Input In [2], in <module>
----> 1 from stable_baselines.common.policies import MlpPolicy

ModuleNotFoundError: No module named 'stable_baselines'

In [3]: from stable_baselines3.common.policies import MlpPolicy
---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Input In [3], in <module>
----> 1 from stable_baselines3.common.policies import MlpPolicy

ImportError: cannot import name 'MlpPolicy' from 'stable_baselines3.common.policies' (/home/raph/RL/stable-baselines/venv/lib/python3.8/site-packages/stable_baselines3/common/policies.py)

To quote the github readme:

Note: Stable-Baselines supports Tensorflow versions from 1.8.0 to 1.14.0. Support for Tensorflow 2 API is planned.

The same github readme also recommends to use stable-baselines3, as stable-baselines is currently only being maintained and its functionality is not extended. Thus, I would not expect the TF1 -> TF2 update any time soon.

If you can not install this version of tensorflow, I suggest to use stable-baselines3 and follow the examples. The code you posted above is not consistent with the stable-baselines3 docs, which import the MlpPolicy as

from stable_baselines3.sac.policies import MlpPolicy