Issue with Multi-GPU Support in IsaacGymEnvs+SKRL

[XInyuSong000]

I encountered an issue while trying to utilize multiple GPUs in SKRL. I modified the config.yaml in isaacgymenvs/cfg as follows:

# device for running physics simulation
sim_device: 'cuda:0'
# device to run RL
rl_device: 'cuda:1'
graphics_device_id: 2

Subsequently, I attempted to run the provided example torch_franka_cube_stack_ppo.py available at SKRL Documentation.

Unfortunately, this resulted in an error:

  File "torch_franka_cube_stack_ppo.py", line 117, in <module>
    trainer.train()
  File "/home/.local/lib/python3.8/site-packages/skrl/trainers/torch/sequential.py", line 77, in train
    self.single_agent_train()
  File "/home/.local/lib/python3.8/site-packages/skrl/trainers/torch/base.py", line 172, in single_agent_train
    actions = self.agents.act(states, timestep=timestep, timesteps=self.timesteps)[0]
  File "/home/.local/lib/python3.8/site-packages/skrl/agents/torch/ppo/ppo.py", line 213, in act
    actions, log_prob, outputs = self.policy.act({"states": self._state_preprocessor(states)}, role="policy")
  File "/home/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/.local/lib/python3.8/site-packages/skrl/resources/preprocessors/torch/running_standard_scaler.py", line 176, in forward
    return self._compute(x, train, inverse)
  File "/home/.local/lib/python3.8/site-packages/skrl/resources/preprocessors/torch/running_standard_scaler.py", line 133, in _compute
    return torch.clamp((x - self.running_mean.float()) / (torch.sqrt(self.running_variance.float()) + self.epsilon),
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!

This leads me to question whether SKRL currently supports training across multiple GPUs. Any guidance or clarification on this matter would be greatly appreciated.

[Toni-SM]

Hi @XInyuSong000

Sorry for late reply

In skrl all components support the specification of a device (through the device argument) where the tensors reside/operate. Such an implementation would allow implementations involving multiple devices... however, I have not implemented the checking and moving of the data for the device configured for each component (something as simple as do data.to(self.device) 😅🙈)