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.