Equivalent for pip install "gymnasium[atari]" in python nix shell

Help
1

Edit:
Just for anyone interested in getting an env running with gymnasium including atari games, I

  1. went to the autorom github copied AutoROM.py and downloaded the roms.
  2. I edited my shellHook to set ALE_ROMS_DIR and also I added dependencies from gymnasium pyproject.toml as was advised in the solution.

Kind of minimal shell.nix for gym env with atari games and pytorch with cuda:

with import <nixpkgs> {
  config = {
    cudaSupport = true;
    allowUnfree = true;
  };
};
let pythonEnv = python312.withPackages(ps: with ps; [
  gymnasium
  ale-py
  torch
]);
in mkShell {
  packages = [
    pythonEnv
  ];
  shellHook =
  ''
    echo "Hello shell"
    export ALE_ROMS_DIR=/path/to/roms/dir/
  '';
}

Original Question:
Hi,
I have to following problem when setting up my shell.nix for my reinforcement learning project.
I can’t figure out how to install the pip package gymnasium correctly. The problem occurs as in a venv I would have to use

    pip install "gymnasium[classic-control]"
    pip install "gymnasium[atari]"
    pip install "gymnasium[accept-rom-licence]"

to get the atari games and accept the corresponding licence. However the nixpkg python312.gymnasium seems to correspond “only” to pip install gymnasium which doesnt include any atari env.
My shell.nix currently looks like this:

with import <nixpkgs> {
  config = {
    cudaSupport = true;
    allowUnfree = true;
  };
};
let pythonEnv = python312.withPackages(ps: with ps; [
  pip
  gymnasium[classic-control]
  gymnasium[atari]
  gymnasium[accept-rom-licence]
  ale-py
  numpy
  torch
  torchvision
  ignite
  tensorboard
  opencv-python
]);
in mkShell {
  packages = [
    pythonEnv
  ];
  shellHook =
  ''
    echo "Hello shell"
  '';
}

However when I check the available pkgs with pip list it looks like it basically only installed gymnasium and in python.

import gymnasium as gym
import ale_py

env = gym.make("ALE/Pong-v5")

throws an error about not being able to find the env.

Has anyone an idea on how to import gymnasium correcly or should I try something else with nix-ld and venv for example.
If someone has a working example for a gym env I would be happy if they could share their approach.

2

Nix will interpret that as gymnasium <array containing the variable classic-control>.

I’m surprised this evaluates, frankly, you should get something like “not a derivation” or “missing variable”.

Nixpkgs doesn’t support this optional dependency shortcut syntax directly, and the package doesn’t add any additional support either. You would need to manually install every package that gymnasium refers to. Alternatively, you could use uv2nix to resolve it instead.

1 Like