From ea21831b6fee2809cc5048338cbf6822d4bf88e8 Mon Sep 17 00:00:00 2001
From: Guanda Li <34904808+Li-Guanda@users.noreply.github.com>
Date: Thu, 2 Feb 2023 18:35:42 +0900
Subject: [PATCH] Update train.py

A command like "python train.py task=Ant headless=True sim_device=cpu rl_device=cpu" can not work correctly.
The reason is "rlg_config_dict" doesn't include the information of "rl_device".

In the "a2c_common.py" of "rl_games", there is a line of code: "self.ppo_device = config.get('device', 'cuda:0')".
So the RL algorithm will always only work on the cuda:0.
---
 isaacgymenvs/train.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/isaacgymenvs/train.py b/isaacgymenvs/train.py
index 783a92d71..3ca00d5ef 100644
--- a/isaacgymenvs/train.py
+++ b/isaacgymenvs/train.py
@@ -140,6 +140,7 @@ def build_runner(algo_observer):
         return runner
 
     rlg_config_dict = omegaconf_to_dict(cfg.train)
+    rlg_config_dict['params']['config']['device']=cfg.rl_device
 
     # convert CLI arguments into dictionory
     # create runner and set the settings