[RLlib] Fix issues with action masking examples

[ArturNiederfahrenhorst]

Why are these changes needed?

Fixes #37707 and a set of other issues.
The action masking example was broken but this was not picked up by CI because the error did not propagate.
The tf1 code path of the example was broken.
The example needed to be liften into RL Modules API.

[kouroshHakha]

LGTM

[ArturNiederfahrenhorst]

@sven1977 Need you tensorflow help here.
It appears that the same code is failing for TF but working for torch.
Here is what I've found out so far:

• Sampling runs fine
• If you compute gradients, they will be infinite (loss produced by learner is not infinite, but gradients are)
• If you don't add the infinite mask (remove the code below), gradients are look normal + no error

inf_mask = tf.maximum(tf.math.log(action_mask), tf.float32.min)
masked_logits = logits + inf_mask

[sven1977]

Let me debug the new example script on tf2. ...

[sven1977]

Root cause: The gradients in tf2 - for some reason - for the last action layer (the one producing the 100 action logits) are being computed as NaNs already in the very first training update step. This then leads to a broken NN and the final crash because one of the action labels is 100 (out of bounds, should be 0 to 99).

[sven1977]

Confirmed that torch is working fine ...

[sven1977]

This is fixed now (new example script is running fine on my end). Waiting for tests to pass, then we can merge.

[sven1977]

LGTM. Thanks for the fix @ArturNiederfahrenhorst !

[PhilippWillms]

@sven1977, @ArturNiederfahrenhorst : Appreciate your work on that topic!

Installing latest nightly build to run the new example file, I stumbled across two topics:

1. Even if I do not call the example with torch, I need to install torch cause in the rl_module example file the corresponding imports are loaded.
2. Which torch version did you use for implementing this PR ? Because with torch==2.0.1, I get the following exception

File "/mnt/d/git/ray/rllib/examples/action_masking.py", line 46, in
from ray.rllib.examples.rl_module.action_masking_rlm import (
File "/home/philipp/miniconda3/envs/tf-39-gpu-linux/lib/python3.9/site-packages/ray/rllib/examples/rl_module/action_masking_rlm.py", line 4, in
from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule
File "/home/philipp/miniconda3/envs/tf-39-gpu-linux/lib/python3.9/site-packages/ray/rllib/algorithms/ppo/torch/ppo_torch_rl_module.py", line 7, in
from ray.rllib.core.rl_module.torch import TorchRLModule
File "/home/philipp/miniconda3/envs/tf-39-gpu-linux/lib/python3.9/site-packages/ray/rllib/core/rl_module/torch/init.py", line 1, in
from .torch_rl_module import TorchRLModule
File "/home/philipp/miniconda3/envs/tf-39-gpu-linux/lib/python3.9/site-packages/ray/rllib/core/rl_module/torch/torch_rl_module.py", line 11, in
from ray.rllib.models.torch.torch_distributions import TorchDistribution
File "/home/philipp/miniconda3/envs/tf-39-gpu-linux/lib/python3.9/site-packages/ray/rllib/models/torch/torch_distributions.py", line 22, in
class TorchDistribution(Distribution, abc.ABC):
File "/home/philipp/miniconda3/envs/tf-39-gpu-linux/lib/python3.9/site-packages/ray/rllib/models/torch/torch_distributions.py", line 51, in TorchDistribution
sample_shape=torch.Size(),

Verification outside of the ray cascade shows interesting behavior of torch leading to the issue

$ python -c "import torch; sample_shape=torch.Size(); print(sample_shape)"

torch.Size([])

[PhilippWillms]

The issue reported earlier does not occur anymore in ray 2.9.1. Tested with Torch 2.0.1 and the examples/action_masking file in the following version: https://github.com/ray-project/ray/blob/84d17cad835665631c2f68f6fb332e2973028fab/rllib/examples/action_masking.py