Debugging script to test rllib initialization with many policies

slow_rllib_init.py

import time

import gym
import ray
from ray.rllib.agents.ppo import PPOTrainer
from ray.rllib.examples.env.multi_agent import MultiAgentCartPole
from ray.tune import register_env

ray.init()

num_policies = 110
num_workers = 10

init_all_at_once = False

# Simple environment with num_policies independent cartpole entities
register_env("multi_agent_cartpole",
             lambda _: MultiAgentCartPole({"num_agents": num_policies}))
single_dummy_env = gym.make("CartPole-v0")
obs_space = single_dummy_env.observation_space
act_space = single_dummy_env.action_space

if init_all_at_once:
    policies = {str(i): (None, obs_space, act_space, {}) for i in range(num_policies)}
else:
    policies = {'0': (None, obs_space, act_space, {})}

policy_mapping_fn = str

print('==================')
print(f'init all at once? {init_all_at_once}')
print(f'num_policies = {num_policies}')
print(f'num_workers {num_workers}')
config = {
    'framework': 'tf',
    'num_workers': num_workers,
    "multiagent": {
        "policies": policies,
        "policy_mapping_fn": policy_mapping_fn,
        "policies_to_train": None,  # All
    }
}
start = time.time()

ppo_trainer = PPOTrainer(
    env="multi_agent_cartpole",
    config=config)

if not init_all_at_once:
    for i in range(1, num_policies):
        _ = ppo_trainer.add_policy(
            policy_id=str(i),
            policy_cls=type(ppo_trainer.get_policy('0'))
        )

ppo_trainer.workers.sync_weights()

print('==================')
print(f"PPO init took {time.time() - start} seconds")
print(f'number of policies created: {len(ppo_trainer.get_weights())}')
print('==================')