[RLlib] [WIP] [MultiAgentEnv Refactor #1] Add new methods to base env

[avnishn]

RLlib Environments are difficult to test. Particularly Base Environments and MultiAgentEnvironments. This is because they are missing necessary required fields in order to test them.

The biggest issue that I face right now is that the MultiAgentEnv API is too loosely defined to write basic environment tests where I would follow this rough workflow:

reset_obs = env.try_reset()
assert env.obs_space_contains(reset_obs)
sampled_actions = env.action_space_sample()
env.send_actions(sampled_actions)
next_obs, dones, rewards, infos = env.poll()

assert env.obs_space_contains(next_obs)
assert isinstance(dones, dict)
assert isinstance(rewards, dict)
assert isinstance(infos, dict)

Because the MultiAgentEnv API is too loosely defined to be able to follow this workflow, I can’t write a BaseEnv checking module.

This PR Adds the necessary BaseEnv methods for making BaseEnvs unit-testable and checkable with an environment checking module.

Checks

• 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 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 :(

[gjoliver]

2 quick questions.

[gjoliver]

wait, is the type still MultiEnvDict here and below??
I thought we are saying gym and multi-agent envs return different types now?

[avnishn]

Yes -- the return type of poll() and try_reset() are MultiEnvDicts, and so I thought it would be appropriate if observations and actions produced by the environment/policy should be able to be easily passed to the environment for checking.

[gjoliver]

ok, it seems to me we are settling on using multi-agent api for single agent env as well, which is totally fine, and probably logical.

[sven1977]

The BaseEnv is never a single-agent env. If there is only one agent and we derive the BaseEnv from e.g. a gym.Env, we auto-create "DUMMY_AGENT_ID" in the env as the agent's ID.

[gjoliver]

a very minor question, do you envision obs/action_space_contains() getting used outside of the environment checking module?

[avnishn]

I can imagine it having other uses for users who are trying to develop their own environments -- I have definitely used functions like this while developing my own environments

[gjoliver]

Looks good to me!