Repository: ray-project/ray
Issue category: Bug Report

[RLlib] action_masking_example.py fails - RLModule build fails with "unexpected keyword argument 'observation_space'"
### What happened + What you expected to happen

Running the `action_masking_rl_module.py` example, which is shipped with 2.39 release, fails at RLModule instantiation.

>   File "C:\Users\Philipp\anaconda3\envs\py311-raynew\Lib\site-packages\ray\rllib\core\rl_module\rl_module.py", line 100, in build
>     module = self.module_class(
>              ^^^^^^^^^^^^^^^^^^
> TypeError: ActionMaskingRLModule.__init__() got an unexpected keyword argument 'observation_space'

I did no change to the file locally. I skipped the CLI args "--enable-new-api-stack", as for PPO the new API stack is enabled by default since of release 2.39

### Versions / Dependencies

python==3.11.9
ray===2.39.0
torch==2.3.1+cu118
gymnasium==1.0.0

### Reproduction script

python ray/rllib/examples/rl_modules/action_masking_rl_module.py

### Issue Severity

Medium: It is a significant difficulty but I can work around it.


Functions to edit:
rllib/examples/rl_modules/classes/action_masking_rlm.py:ActionMaskingRLModule.__init__
rllib/examples/rl_modules/classes/action_masking_rlm.py:ActionMaskingTorchRLModule.compute_values

Patch:
diff --git a/rllib/examples/rl_modules/classes/action_masking_rlm.py b/rllib/examples/rl_modules/classes/action_masking_rlm.py
index 992802ebb13a..626554a6434c 100644
--- a/rllib/examples/rl_modules/classes/action_masking_rlm.py
+++ b/rllib/examples/rl_modules/classes/action_masking_rlm.py
@@ -1,10 +1,11 @@
 import gymnasium as gym
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple, Union
 
 from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule
 from ray.rllib.core.columns import Columns
 from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI
-from ray.rllib.core.rl_module.rl_module import RLModule, RLModuleConfig
+from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
+from ray.rllib.core.rl_module.rl_module import RLModule
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.framework import try_import_torch
 from ray.rllib.utils.torch_utils import FLOAT_MIN
@@ -32,9 +33,17 @@ class ActionMaskingRLModule(RLModule):
     """
 
     @override(RLModule)
-    def __init__(self, config: RLModuleConfig):
+    def __init__(
+        self,
+        observation_space: Optional[gym.Space] = None,
+        action_space: Optional[gym.Space] = None,
+        inference_only: Optional[bool] = None,
+        learner_only: bool = False,
+        model_config: Optional[Union[dict, DefaultModelConfig]] = None,
+        catalog_class=None,
+    ):
         # If observation space is not of type `Dict` raise an error.
-        if not isinstance(config.observation_space, gym.spaces.dict.Dict):
+        if not isinstance(observation_space, gym.spaces.dict.Dict):
             raise ValueError(
                 "This RLModule requires the environment to provide a "
                 "`gym.spaces.Dict` observation space of the form: \n"
@@ -46,15 +55,22 @@ def __init__(self, config: RLModuleConfig):
         # the action mask and the original observation space, the 'RLModule'
         # receives only the `"observation"` element of the space, but not the
         # action mask.
-        self.observation_space_with_mask = config.observation_space
-        config.observation_space = config.observation_space["observations"]
+        self.observation_space_with_mask = observation_space
+        self.observation_space = observation_space["observations"]
 
         # Keeps track if observation specs have been checked already.
         self._checked_observations = False
 
         # The PPORLModule, in its constructor will build networks for the original
         # observation space (i.e. without the action mask).
-        super().__init__(config)
+        super().__init__(
+            observation_space=self.observation_space,
+            action_space=action_space,
+            inference_only=inference_only,
+            learner_only=learner_only,
+            model_config=model_config,
+            catalog_class=catalog_class,
+        )
 
 
 class ActionMaskingTorchRLModule(ActionMaskingRLModule, PPOTorchRLModule):
@@ -100,11 +116,13 @@ def _forward_train(
 
     @override(ValueFunctionAPI)
     def compute_values(self, batch: Dict[str, TensorType], embeddings=None):
-        # Preprocess the batch to extract the `observations` to `Columns.OBS`.
-        action_mask, batch = self._preprocess_batch(batch)
-        # NOTE: Because we manipulate the batch we need to add the `action_mask`
-        # to the batch to access them in `_forward_train`.
-        batch["action_mask"] = action_mask
+        # Check, if the observations are still in `dict` form.
+        if isinstance(batch[Columns.OBS], dict):
+            # Preprocess the batch to extract the `observations` to `Columns.OBS`.
+            action_mask, batch = self._preprocess_batch(batch)
+            # NOTE: Because we manipulate the batch we need to add the `action_mask`
+            # to the batch to access them in `_forward_train`.
+            batch["action_mask"] = action_mask
         # Call the super's method to compute values for GAE.
         return super().compute_values(batch, embeddings)