Some questions when I run distral code

[c4cld]

Description

What issue are you facing?

When I run the following code which you mention in https://mtrl.readthedocs.io/en/latest/pages/tutorials/baseline.html, error occurs.

MUJOCO_GL="osmesa" LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu/nvidia-opengl/:$LD_LIBRARY_PATH PYTHONPATH=. python3 -u main.py
setup=hipbmdp
env=dmcontrol-finger-spin-distribution-v0
agent=distral
setup.seed=1
agent.distral_alpha=1.0
agent.distral_beta=1.0
replay_buffer.batch_size=256

How to reproduce

Add a minimal example to reproduce the issue.
MUJOCO_GL="osmesa" LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu/nvidia-opengl/:$LD_LIBRARY_PATH PYTHONPATH=. python3 -u main.py
setup=hipbmdp
env=dmcontrol-finger-spin-distribution-v0
agent=distral
setup.seed=1
agent.distral_alpha=1.0
agent.distral_beta=1.0
replay_buffer.batch_size=256

Stack trace / error message

Paste the stack trace/error message to Gist
and paste the link here.

The link of error message
https://gist.github.com/ZJP-c4cld/a0e052ee2d4ca8c88044ca7062fc3ab2

System information

• MTRL Version : The output of "pip list" in mtrl conda environment is as follows:
• absl-py 0.12.0
  aiohttp 3.7.4.post0
  aiohttp-cors 0.7.0
  aioredis 1.3.1
  alabaster 0.7.12
  antlr4-python3-runtime 4.8
  anyio 3.1.0
  apipkg 1.5
  appdirs 1.4.4
  argcomplete 1.12.3
  argon2-cffi 20.1.0
  astor 0.8.1
  async-generator 1.10
  async-timeout 3.0.1
  attrs 21.2.0
  Babel 2.9.1
  backcall 0.2.0
  beautifulsoup4 4.9.3
  black 20.8b1
  bleach 3.3.0
  blessings 1.7
  Brotli 1.0.9
  cached-property 1.5.2
  cachetools 4.2.2
  certifi 2020.12.5
  cffi 1.14.5
  cfgv 3.3.0
  chardet 4.0.0
  click 8.0.0
  cloudpickle 1.2.2
  colorama 0.4.4
  colorful 0.5.4
  colorlog 4.8.0
  contextvars 2.4
  cycler 0.10.0
  Cython 0.29.23
  dataclasses 0.8
  decorator 5.0.9
  defusedxml 0.7.1
  distlib 0.3.2
  dm-tree 0.1.6
  docopt 0.6.2
  docutils 0.17.1
  entrypoints 0.3
  execnet 1.8.1
  fasteners 0.16
  filelock 3.0.12
  flake8 3.8.4
  flake8-bugbear 20.11.1
  flake8-comprehensions 3.3.1
  flake8-docstrings 1.5.0
  Flask 2.0.1
  Flask-Compress 1.9.0
  flatbuffers 2.0
  funcsigs 1.0.2
  future 0.18.2
  gast 0.4.0
  glfw 2.1.0
  google 3.0.0
  google-api-core 1.27.0
  google-auth 1.30.0
  google-auth-oauthlib 0.4.4
  googleapis-common-protos 1.53.0
  gpustat 0.6.0
  grpcio 1.37.1
  gym 0.18.3
  h5py 3.1.0
  highway-env 1.0.dev0
  hiplot 0.1.22
  hiredis 2.0.0
  Hydra 2.5
  hydra-core 1.0.5
  hydra-submitit-launcher 1.0.1
  identify 2.2.6
  idna 2.10
  idna-ssl 1.1.0
  imageio 2.9.0
  imageio-ffmpeg 0.4.3
  imagesize 1.2.0
  immutables 0.15
  importlib-metadata 4.0.1
  importlib-resources 5.1.4
  incremental 21.3.0
  iniconfig 1.1.1
  ipykernel 5.5.5
  ipython 7.16.1
  ipython-genutils 0.2.0
  ipywidgets 7.6.3
  isort 5.7.0
  itsdangerous 2.0.1
  jedi 0.18.0
  Jinja2 3.0.1
  joblib 1.0.1
  json5 0.9.5
  jsonschema 3.2.0
  jupyter 1.0.0
  jupyter-client 6.1.12
  jupyter-console 6.4.0
  jupyter-contrib-core 0.3.3
  jupyter-core 4.7.1
  jupyter-nbextensions-configurator 0.4.1
  jupyter-server 1.8.0
  jupyterlab 3.0.3
  jupyterlab-pygments 0.1.2
  jupyterlab-server 2.5.2
  jupyterlab-widgets 1.0.0
  Keras-Applications 1.0.8
  Keras-Preprocessing 1.1.2
  kiwisolver 1.3.1
  latexcodec 2.0.1
  llvmlite 0.36.0
  lockfile 0.12.2
  lz4 3.1.3
  Markdown 3.3.4
  MarkupSafe 2.0.1
  matplotlib 3.3.3
  mccabe 0.6.1
  mistune 0.8.4
  mllogger 0.7
  msgpack 1.0.2
  mtenv 1.0 /media/dl_e/17EF48822800F82E/zjp/code/procgen/mtrl/src/mtenv
  mujoco-py 2.0.2.13
  multidict 5.1.0
  mypy 0.790
  mypy-extensions 0.4.3
  nbclassic 0.3.1
  nbclient 0.5.3
  nbconvert 6.0.7
  nbformat 5.1.3
  nest-asyncio 1.5.1
  nodeenv 1.6.0
  notebook 6.4.0
  nox 2020.12.31
  numba 0.53.1
  numpy 1.19.5
  nvidia-ml-py3 7.352.0
  oauthlib 3.1.0
  omegaconf 2.0.6
  opencensus 0.7.13
  opencensus-context 0.1.2
  opencv-python 4.5.1.48
  packaging 20.9
  pandas 1.1.5
  pandocfilters 1.4.3
  parso 0.8.2
  pathspec 0.8.1
  pexpect 4.8.0
  pickleshare 0.7.5
  Pillow 7.1.0
  pip 21.1.1
  pluggy 0.13.1
  pockets 0.9.1
  pre-commit 2.9.3
  prometheus-client 0.10.1
  prompt-toolkit 3.0.18
  protobuf 3.17.0
  psutil 5.8.0
  ptyprocess 0.7.0
  py 1.10.0
  py-spy 0.3.7
  pyasn1 0.4.8
  pyasn1-modules 0.2.8
  pybtex 0.24.0
  pybtex-docutils 1.0.0
  pycodestyle 2.6.0
  pycparser 2.20
  pydocstyle 6.1.1
  pyflakes 2.2.0
  pygame 1.9.6
  pyglet 1.5.0
  Pygments 2.9.0
  pymongo 3.11.2
  pynndescent 0.5.2
  pyparsing 2.4.7
  pyrsistent 0.17.3
  pytest 6.2.1
  pytest-forked 1.3.0
  pytest-xdist 2.2.0
  python-dateutil 2.8.1
  pytz 2020.1
  PyYAML 5.4.1
  pyzmq 22.1.0
  qtconsole 5.1.0
  QtPy 1.9.0
  ray 1.0.1.post1
  redis 3.4.1
  regex 2021.4.4
  requests 2.25.1
  requests-oauthlib 1.3.0
  rl-agents 1.0.dev0
  rsa 4.7.2
  scikit-learn 0.24.2
  scipy 1.5.4
  seaborn 0.11.0
  Send2Trash 1.5.0
  setproctitle 1.2.2
  setuptools 51.1.1
  six 1.16.0
  sniffio 1.2.0
  snowballstemmer 2.1.0
  soupsieve 2.2.1
  Sphinx 3.4.3
  sphinx-autodoc-annotation 1.0.post1
  sphinx-copybutton 0.3.1
  sphinx-rtd-theme 0.5.1
  sphinxcontrib-applehelp 1.0.2
  sphinxcontrib-bibtex 2.1.3
  sphinxcontrib-devhelp 1.0.2
  sphinxcontrib-htmlhelp 2.0.0
  sphinxcontrib-jsmath 1.0.1
  sphinxcontrib-napoleon 0.7
  sphinxcontrib-qthelp 1.0.3
  sphinxcontrib-serializinghtml 1.1.5
  submitit 1.3.3
  svgwrite 1.4.1
  tabulate 0.8.9
  tensorboard 2.4.0
  tensorboard-plugin-wit 1.8.0
  tensorboardX 2.1
  tensorflow 1.12.0
  termcolor 1.1.0
  terminado 0.10.0
  testpath 0.5.0
  threadpoolctl 2.1.0
  toml 0.10.2
  torch 1.7.1+cu110
  torchvision 0.8.2+cu110
  tornado 6.1
  towncrier 19.2.0
  traitlets 4.3.3
  Tree 0.2.4
  typed-ast 1.4.3
  typing 3.7.4.3
  typing-extensions 3.7.4.3
  umap-learn 0.5.1
  urllib3 1.26.4
  virtualenv 20.4.7
  wcwidth 0.2.5
  webencodings 0.5.1
  websocket-client 1.0.1
  Werkzeug 2.0.1
  wheel 0.36.2
  widgetsnbextension 3.5.1
  yamllint 1.25.0
  yarl 1.6.3
  zipp 3.4.1
• MTRL environment Name : mtrl
• Python version : 3.6.5

Any other information

Add any other information here.

[shagunsodhani]

Hi! thanks for raising this issue. I have added a fix in #7. I will merge once the tests pass.

[c4cld]

Thank you for your help ^_^