[RLlib] Fix env_check for parametric actions (with action mask)

[inpefess]

Why are these changes needed?

Ray RLlib has a great example of using a parametric actions environment, but now it works only with self._skip_env_checking = True. Gymnasium action spaces have a mask argument to their sample method. We apply this feature to fix the env_check behaviour in the parametric actions environments case.

Related issue number

Closes #23925

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

Discussion

1. I added a test to rllib/utils/tests/test_env_check.py because it seems to belong there, but seven tests from this module fail in the master branch
2. The assumption that having a dictionary with the "action_mask" key is a standard parametric actions implementation seems too strong (one can use other names or add the mask to info, for example). If it's a standard or recommended way to do that in Ray RLlib, then one should also mention it in the documentation (now it's not mentioned at all, although the previous version of the documentation page was a bit more verbose). It seems to be a "standard" way, for example, in AlphaZero it's the same and in a couple of other examples.
3. Apart from ParametricActionsCartPole using the action_mask key, there is a ParametricActionsCartPoleNoEmbeddings having a key with the same meaning but called valid_avail_actions_mask. I've changed it to action_mask

[inpefess]

@avnishn Could you review this PR, please? Failing checks don't seem to be related to the changes proposed.

[avnishn]

The assumption that having a dictionary with the "action_mask" key is a standard parametric actions implementation seems too strong (one can use other names or add the mask to info, for example). If it's a standard or recommended way to do that in Ray RLlib, then one should also mention it in the documentation (now it's not mentioned at all, although the previous version of the documentation page was a bit more verbose). It seems to be a "standard" way, for example, in AlphaZero it's the same and in a couple of other examples.

I think that is very much the case. We have to think about if we want to make an opinionated design decision on how to support action masking in RLlib.

On an initial thought, I think that we'd rather not make an opinionated design decision here. I'll talk with @sven1977 who probably has more opinions on this and get back to you.

[sven1977]

Very nice!

[sven1977]

This is a really cool feature of gymnasium, actually :) which I didn't know about.

[sven1977]

Awesome PR. Thanks for filing this @inpefess . I would like to suggest one enhancement to make the "action_mask" key not hard-coded.
Can we add an additional config setting in therllib/algorithms/algorithm_config.py::AlgorithmConfig::environment() method to be able to customize the actual value of this action mask key in the observation_space?

something along the lines of:

config.environment("my_env", env_config=..., action_mask_key="action_mask")

Set the default value of self.action_mask_key = "action_mask" in the AlgorithmConfig c'tor.

Then use that value instead of the hard-coded one in the pre-check.
You might have to change the signature of check_env to pass along the AlgorithmConfig (from within the RolloutWorker) so that it has access to that configuration.

[sven1977]

Awesome, thanks for adding this so quickly! I think it's ready to merge now. Just waiting for tests to finish ..

[sven1977]

Looks great now. Thanks again @inpefess !

[afennelly-mitre]

@inpefess @sven1977 @avnishn I was browsing through this PR from earlier this year, and wanted to verify if my assumption is correct about the changes in this PR:

the env_check for parametric actions (with action mask) will only work if the underlying environment is a gym.Env, and will not work if the environment is say, a VectorEnv or MultiAgentEnv since the logic is only added to the check_gym_environments() method?