DLR-RM / DLR-RM/stable-baselines3
[Bug]: SubprocVecEnv ignores specified CUDA device and uses GPU 0
- Dominant language
- Python
- Stars
- 13.8k
- Forks
- 2.2k
- Avg merge
- 1h 35m
- Merged PRs (30d)
- 2
Description
### 🐛 Bug
## Description
When using `SubprocVecEnv` with multiple environments, all subprocesses ignore the GPU device specified for the main process and default to GPU 0, regardless of which GPU was set with `torch.cuda.set_device()`.
## Reproduction
I created a minimal reproduction script that clearly shows the issue:
```python
import os
import torch
import numpy as np
import gymnasium as gym
from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv
class GPUTestEnv(gym.Env):
def __init__(self, env_id=0):
super().__init__()
self.observation_space = gym.spaces.Box(low=-1, high=1, shape=(4,), dtype=np.float32)
self.action_space = gym.spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32)
self.env_id = env_id
if torch.cuda.is_available():
device = torch.cuda.current_device()
device_name = torch.cuda.get_device_name(device)
print(f"Env {self.env_id} created on GPU {device} ({device_name}) - PID: {os.getpid()}")
def reset(self, seed=None, options=None):
if torch.cuda.is_available():
test_tensor = torch.ones(1, device="cuda")
device_id = test_tensor.device.index
print(f"Env {self.env_id} - reset() using GPU: {device_id}")
return np.zeros(4, dtype=np.float32), {}
def step(self, action):
if torch.cuda.is_available():
test_tensor = torch.ones(1, device="cuda")
device_id = test_tensor.device.index
print(f"Env {self.env_id} - step() using GPU: {device_id}")
return np.zeros(4, dtype=np.float32), 0.0, False, False, {}
# Specify GPU 3
torch.cuda.set_device(3)
print(f"Main process current device: {torch.cuda.current_device()}")
# Create environments
env_fns = [lambda idx=i: GPUTestEnv(idx) for i in range(3)]
# DummyVecEnv correctly uses GPU 3
print("\n----- Testing DummyVecEnv -----")
dummy_env = DummyVecEnv(env_fns)
dummy_env.reset()
dummy_env.step(np.zeros((3, 2)))
dummy_env.close()
# SubprocVecEnv incorrectly uses GPU 0
print("\n----- Testing SubprocVecEnv -----")
subproc_env = SubprocVecEnv(env_fns, start_method="spawn")
subproc_env.reset()
subproc_env.step(np.zeros((3, 2)))
subproc_env.close()
```
## Output
```
Main process current device: 3
----- Testing DummyVecEnv -----
Env 0 created on GPU 3 (NVIDIA A100 80GB PCIe) - PID: 1795905
Env 1 created on GPU 3 (NVIDIA A100 80GB PCIe) - PID: 1795905
Env 2 created on GPU 3 (NVIDIA A100 80GB PCIe) - PID: 1795905
Env 0 - reset() using GPU: 3
Env 1 - reset() using GPU: 3
Env 2 - reset() using GPU: 3
Env 0 - step() using GPU: 3
Env 1 - step() using GPU: 3
Env 2 - step() using GPU: 3
----- Testing SubprocVecEnv -----
Env 0 created on GPU 0 (NVIDIA A100 80GB PCIe) - PID: 1796000
Env 0 - reset() using GPU: 0
Env 0 - step() using GPU: 0
Env 2 created on GPU 0 (NVIDIA A100 80GB PCIe) - PID: 1796002
Env 2 - reset() using GPU: 0
Env 2 - step() using GPU: 0
Env 1 created on GPU 0 (NVIDIA A100 80GB PCIe) - PID: 1796001
Env 1 - reset() using GPU: 0
Env 1 - step() using GPU: 0
```
## Environment
- Stable Baselines 3 version: 2.5.0
- PyTorch: 2.6.0+cu124
- CUDA: 12.4
- GPUs: 4x NVIDIA A100 80GB PCIe
- gymnasium:1.0.0
### Checklist
- [x] My issue does not relate to a custom gym environment. (Use the custom gym env template instead)
- [x] I have checked that there is no similar [issue](https://github.com/DLR-RM/stable-baselines3/issues) in the repo
- [x] I have read the [documentation](https://stable-baselines3.readthedocs.io/en/master/)
- [x] I have provided a [minimal and working](https://github.com/DLR-RM/stable-baselines3/issues/982#issuecomment-1197044014) example to reproduce the bug
- [x] I've used the [markdown code blocks](https://help.github.com/en/articles/creating-and-highlighting-code-blocks) for both code and stack traces.
Contributor guide
Research direction
Start at the SubprocVecEnv entry point and compare its spawned-process behavior with DummyVecEnv, using the provided reproduction with torch.cuda.set_device(3). Run the reproduction and verify that subprocess environment creation, reset, and step use the selected device rather than GPU 0.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100