"The observation_space and action_space were not given, can't verify new environments" when loading expert policy

[boshi-an]

Bug description

KeyError: The observation_space and action_space were not given, can't verify new environments.

Steps to reproduce

The code to reproduce is from Imitation's official documents. (https://imitation.readthedocs.io/en/latest/tutorials/1_train_bc.html)

import numpy as np
import gymnasium as gym
from stable_baselines3.common.evaluation import evaluate_policy

from imitation.algorithms import bc
from imitation.data import rollout
from imitation.data.wrappers import RolloutInfoWrapper
from imitation.policies.serialize import load_policy
from imitation.util.util import make_vec_env

rng = np.random.default_rng(0)
env = make_vec_env(
    "seals:seals/CartPole-v0",
    rng=rng,
    n_envs=1,
    post_wrappers=[lambda env, _: RolloutInfoWrapper(env)],  # for computing rollouts
)
expert = load_policy(
    "ppo-huggingface",
    organization="HumanCompatibleAI",
    env_name="seals-CartPole-v0",
    venv=env,
    # observation_space=env.observation_space,
    # action_space=env.action_space
)

Environment

• Operating system and version:
• Python version:
• Output of pip freeze --all:

absl-py==1.4.0
aiohttp==3.9.5
aiosignal==1.3.1
alembic==1.13.2
asttokens==2.2.1
async-timeout==4.0.3
attrs==23.2.0
backcall @ file:///home/conda/feedstock_root/build_artifacts/backcall_1592338393461/work
cachetools==5.3.1
certifi==2023.5.7
cffi==1.15.1
charset-normalizer==3.1.0
click==8.1.3
cloudpickle==3.0.0
cmake==3.26.4
colorama==0.4.6
colorlog==6.8.2
comm @ file:///home/conda/feedstock_root/build_artifacts/comm_1710320294760/work
contourpy==1.1.0
cycler==0.11.0
Cython==0.29.36
datasets==2.20.0
debugpy @ file:///home/conda/feedstock_root/build_artifacts/debugpy_1707444446795/work
decorator @ file:///home/conda/feedstock_root/build_artifacts/decorator_1641555617451/work
dill==0.3.8
docker-pycreds==0.4.0
docopt==0.6.2
executing==1.2.0
Farama-Notifications==0.0.4
fasteners==0.15
filelock==3.12.2
flatten-dict==0.4.2
fonttools==4.40.0
free-mujoco-py==2.1.6
frozenlist==1.4.1
fsspec==2024.5.0
gitdb==4.0.11
GitPython==3.1.43
glfw==1.12.0
google-auth==2.21.0
google-auth-oauthlib==0.4.6
greenlet==3.0.3
grpcio==1.56.0
gym==0.13.0
gymnasium==0.29.1
h5py==3.7.0
huggingface-hub==0.23.4
huggingface-sb3==3.0
idna==3.4
imageio==2.31.1
imitation==1.0.0
importlib-metadata==6.7.0
importlib-resources==5.12.0
ipdb==0.13.13
ipykernel @ file:///home/conda/feedstock_root/build_artifacts/ipykernel_1717717528849/work
ipython @ file:///home/conda/feedstock_root/build_artifacts/ipython_1680185408135/work
# Editable install with no version control (isaacgym==1.0rc4)
-e /home/boshi/Documents/Packages/isaacgym4/python
jax-jumpy==1.0.0
jedi==0.18.2
Jinja2==3.1.2
joblib==1.4.2
jsonpickle==3.2.2
jupyter_client @ file:///home/conda/feedstock_root/build_artifacts/jupyter_client_1716472197302/work
jupyter_core @ file:///home/conda/feedstock_root/build_artifacts/jupyter_core_1710257397447/work
kiwisolver==1.4.4
lit==16.0.6
Mako==1.3.5
Markdown==3.4.3
markdown-it-py==3.0.0
MarkupSafe==2.1.3
matplotlib==3.7.1
matplotlib-inline==0.1.6
mdurl==0.1.2
monotonic==1.6
mpmath==1.3.0
multidict==6.0.5
multiprocess==0.70.16
munch==4.0.0
MyoSuite==1.6.1
nest_asyncio @ file:///home/conda/feedstock_root/build_artifacts/nest-asyncio_1705850609492/work
networkx==3.1
ninja==1.11.1
numpy==1.22.4
nvidia-cublas-cu11==11.10.3.66
nvidia-cuda-cupti-cu11==11.7.101
nvidia-cuda-nvrtc-cu11==11.7.99
nvidia-cuda-runtime-cu11==11.7.99
nvidia-cudnn-cu11==8.5.0.96
nvidia-cufft-cu11==10.9.0.58
nvidia-curand-cu11==10.2.10.91
nvidia-cusolver-cu11==11.4.0.1
nvidia-cusparse-cu11==11.7.4.91
nvidia-nccl-cu11==2.14.3
nvidia-nvtx-cu11==11.7.91
oauthlib==3.2.2
optuna==3.6.1
packaging==23.1
pandas==2.0.3
parso==0.8.3
patchelf==0.17.2.1
pexpect==4.8.0
pickleshare @ file:///home/conda/feedstock_root/build_artifacts/pickleshare_1602536217715/work
Pillow==10.0.0
pip==23.1.2
platformdirs @ file:///home/conda/feedstock_root/build_artifacts/platformdirs_1715777629804/work
prompt-toolkit==3.0.39
protobuf==3.19.6
psutil @ file:///home/conda/feedstock_root/build_artifacts/psutil_1719274595110/work
ptyprocess @ file:///home/conda/feedstock_root/build_artifacts/ptyprocess_1609419310487/work/dist/ptyprocess-0.7.0-py2.py3-none-any.whl
pure-eval @ file:///home/conda/feedstock_root/build_artifacts/pure_eval_1642875951954/work
py-cpuinfo==9.0.0
py-spy==0.3.14
pyarrow==16.1.0
pyarrow-hotfix==0.6
pyasn1==0.5.0
pyasn1-modules==0.3.0
pybullet==3.2.6
pycparser==2.21
pygame==2.6.0
pyglet==2.0.8
Pygments==2.15.1
pyparsing==3.1.0
python-dateutil==2.8.2
pytz==2023.3
PyYAML==6.0
pyzmq @ file:///home/conda/feedstock_root/build_artifacts/pyzmq_1715024398647/work
requests==2.32.3
requests-oauthlib==1.3.1
rich==13.7.1
rsa==4.9
sacred==0.8.5
sb3-contrib==1.6.0
scikit-learn==1.3.2
scipy==1.10.1
seals==0.2.1
sentry-sdk==2.7.1
setproctitle==1.3.3
setuptools==67.8.0
six @ file:///home/conda/feedstock_root/build_artifacts/six_1620240208055/work
sk-video==1.1.10
smmap==5.0.1
SQLAlchemy==2.0.31
stable_baselines3==2.3.2
stack-data @ file:///home/conda/feedstock_root/build_artifacts/stack_data_1669632077133/work
sympy==1.12
tensorboard==2.10.0
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.1
termcolor==2.3.0
threadpoolctl==3.5.0
tomli==2.0.1
torch==2.0.1
torchvision==0.15.2
tornado @ file:///home/conda/feedstock_root/build_artifacts/tornado_1717722826518/work
tqdm==4.66.4
traitlets @ file:///home/conda/feedstock_root/build_artifacts/traitlets_1713535121073/work
triton==2.0.0
typing_extensions @ file:///home/conda/feedstock_root/build_artifacts/typing_extensions_1717802530399/work
tzdata==2023.3
urllib3==1.26.16
wandb==0.17.3
wasabi==1.1.3
wcwidth @ file:///home/conda/feedstock_root/build_artifacts/wcwidth_1704731205417/work
Werkzeug==2.3.6
wheel==0.38.4
wrapt==1.16.0
xxhash==3.4.1
yappi==1.4.0
yarl==1.9.4
zipp==3.15.0

[chrimenz]

same here. partial output when trying to download pretrained expert inside a notebook:

File ~/miniforge3/lib/python3.9/site-packages/imitation/policies/serialize.py:71, in load_stable_baselines_model(cls, path, venv, **kwargs)
65 if vec_normalize_path.exists():
66 raise FileExistsError(
67 "Outdated policy format: we do not support restoring normalization "
68 f"statistics from '{vec_normalize_path}'",
69 )
---> 71 return cls.load(path_obj, env=venv, **kwargs)

File ~/miniforge3/lib/python3.9/site-packages/stable_baselines3/common/base_class.py:704, in BaseAlgorithm.load(cls, path, env, device, custom_objects, print_system_info, force_reset, **kwargs)
698 raise ValueError(
699 f"The specified policy kwargs do not equal the stored policy kwargs."
700 f"Stored kwargs: {data['policy_kwargs']}, specified kwargs: {kwargs['policy_kwargs']}"
701 )
703 if "observation_space" not in data or "action_space" not in data:
--> 704 raise KeyError("The observation_space and action_space were not given, can't verify new environments")
706 # Gym -> Gymnasium space conversion
707 for key in {"observation_space", "action_space"}:

KeyError: "The observation_space and action_space were not given, can't verify new environments"

[zyegao]

I solved this problem by the following method：
Change the last line of code in the function "load_stable_baselines_model" under the file “imitation/src/imitation/policies/serialize.py” to

return cls.load(path_obj, env=venv, custom_objects = {'observation_space': venv.observation_space, 'action_space': venv.action_space}, **kwargs)

This code originally was：

return cls.load(path_obj, env=venv, **kwargs)