20ef4a860325d1735df038d4e4b8c7de5e013ed3,rllib/tests/test_avail_actions_qmix.py,,,#,43

Before Change

        return obs, rewards, dones, {}


if __name__ == "__main__":
    grouping = {
        "group_1": ["agent_1"],  // trivial grouping for testing
    }
    obs_space = Tuple([AvailActionsTestEnv.observation_space])
    act_space = Tuple([AvailActionsTestEnv.action_space])
    register_env(
        "action_mask_test",
        lambda config: AvailActionsTestEnv(config).with_agent_groups(
            grouping, obs_space=obs_space, act_space=act_space))

    ray.init()
    agent = QMixTrainer(
        env="action_mask_test",
        config={
            "num_envs_per_worker": 5,  // test with vectorization on
            "env_config": {
                "avail_action": 3,
            },
        })
    for _ in range(5):
        agent.train()  // OK if it doesn"t trip the action assertion error
    assert agent.train()["episode_reward_mean"] == 21.0

After Change

        assert agent.train()["episode_reward_mean"] == 21.0


if __name__ == "__main__":
    import pytest
    import sys
    sys.exit(pytest.main(["-v", __file__]))

Instances