[RLlib] Issue 28428: QMix not working w/ GPUs. (ray-project#31299)

Signed-off-by: tmynn <hovhannes.tamoyan@gmail.com>
ju2ez · Jan 25, 2023 · 933a07a · 933a07a
1 parent 90d532a
commit 933a07a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/rllib/algorithms/qmix/qmix.py b/rllib/algorithms/qmix/qmix.py
@@ -79,6 +79,12 @@ def __init__(self):
         self.optim_eps = 0.00001
         self.grad_clip = 10
 
+        # QMix-torch overrides the TorchPolicy's learn_on_batch w/o specifying a
+        # alternative `learn_on_loaded_batch` alternative for the GPU.
+        # TODO: This hack will be resolved once we move all algorithms to the new
+        #  RLModule/RLTrainer APIs.
+        self.simple_optimizer = True
+
         # Override some of AlgorithmConfig's default values with QMix-specific values.
         # .training()
         self.lr = 0.0005

diff --git a/rllib/policy/torch_policy_v2.py b/rllib/policy/torch_policy_v2.py
@@ -182,7 +182,7 @@ def __init__(
 
         # Backward compatibility workaround so Policy will call self.loss() directly.
         # TODO(jungong): clean up after all policies are migrated to new sub-class
-        # implementation.
+        #  implementation.
         self._loss = None
 
         # Store, which params (by index within the model's list of