Package: pax-rl
Pax: Meta|Multi Agent Learning in JAXPax is an experiment runner for multi and meta agent research built on top of JAX. It supports "other agent shaping", "multi agent RL" and "single agent RL" experiments. It supports evolutionary and RL-based optimisation.Pax (noun) - a period of peace that has been forced on a large area, such as an empire or even the whole worldPax is composed of 3 components: Environments, Agents and Runners.EnvironmentsEnvironments are similar togymnax.frompax.envs.iterated_matrix_gameimport(IteratedMatrixGame,EnvParams,)env=IteratedMatrixGame(num_inner_steps=5)env_params=EnvParams(payoff_matrix=payoff)# 0 = Defect, 1 = Cooperateactions=(jnp.ones(()),jnp.ones(()))obs,env_state=env.reset(rng,env_params)done=Falsewhilenotdone:obs,env_state,rewards,done,info=env.step(rng,env_state,actions,env_params)We can compose these with JAX built-in functinsjit,vmap,pmapandlax.scan.importIteratedMatrixGame,EnvParamsimportjax.numpyasjnp# batch over env initalisationsnum_envs=2payoff=[[2,2],[0,3],[3,0],[1,1]]rollout_length=50rng=jnp.concatenate([jax.random.PRNGKey(0),jax.random.PRNGKey(1)]).reshape(num_envs,-1)env=IteratedMatrixGame(num_inner_steps=rollout_length)env_params=EnvParams(payoff_matrix=payoff)action=jnp.ones((num_envs,),dtype=jnp.float32)# we want to batch over rngs, actionsenv.step=jax.vmap(env.step,in_axes=(0,None,0,None),out_axes=(0,None,0,0,0),)obs,env_state=env.reset(rng,env_params)# lets scan the rollout for speeddefrollout(carry,unused):carry=(_,env_state,env_rng)actions=(action,action)obs,env_state,rewards,done,info=env.step(env_rng,env_state,actions,env_params)return(obs,env_state,env_rng),(obs,actions,rewards,done,)final_state,trajectory=jax.lax.scan(rollout,(obs,env_state,rng),rollout_length)AgentsThe agent interface is as follows:importjax.numpyasjnpimportAgentargs={"hidden":16,"observation_spec":5}rng=jax.random.PRNGKey(0)bs=1init_hidden=jnp.zeros((bs,args.hidden))obs=jnp.ones((bs,5))agent=Agent(args)state,memory=agent.make_initial_state(rng,init_hidden)action,state,mem=agent.policy(rng,obs,mem)state,memory,stats=agent.update(traj_batch,obs,state,mem)mem=agent.reset_memory(mem,False)Note thatmake_initial_state,policy,updateandreset_memoryall supportjit,vmapandlax.scan. Allowing you to compile more of your experiment toXLA.# batch MemoryState not TrainingStateagent.batch_reset=jax.jit(jax.vmap(agent.reset_memory,(0,None),0),static_argnums=1,)agent.batch_policy=jax.jit(jax.vmap(agent._policy,(None,0,0),(0,None,0)))agent1.batch_init=jax.vmap(agent.make_initial_state,(None,0),(None,0),)RunnersWe can finally combine all the above into our runner code. This is where you'd expect to write most custom logic for your own experimental set up,def_rollout(carry,unused):"""Runner for inner episode"""(rngs,obs,a_state,a_mem,env_state,env_params,)=carry# unpack rngsrngs=self.split(rngs,4)action,a_state,new_a_mem=agent1.batch_policy(a_state,obs[0],a_mem,)next_obs,env_state,rewards,done,info=env.step(rngs,env_state,(action,action),env_params,)traj=Sample(obs1,action,rewards[0],new_a1_mem.extras["log_probs"],new_a1_mem.extras["values"],done,a1_mem.hidden,)return(rngs,next_obs,a1_state,new_a1_mem,env_state,env_params,),(traj1,traj2,)agent=Agent(args)state,memory=agent.make_initial_state(rng,init_hidden)for_inrange(num_updates):final_timestep,batch_trajectory=jax.lax.scan(_rollout,((obs,env_state,rng),rollout_length),10,)_,obs,rewards,a1_state,a1_mem,_,_=final_timestepstate,memory,stats=agent.update(batch_trajectory,obs[0],state,memory)Note this isn't even a fully optimised example - we could jit the outer loop!InstallationPax is written in pure Python, but depends on C++ code via JAX.Because JAX installation is different depending on your CUDA version, Haiku does not list JAX as a dependency in requirements.txt.First, follow these instructions to install JAX with the relevant accelerator support.General InformationThe project entrypoint ispax/experiment.py. The simplest command to run a game would be:python-mpax.experimentWe currently useWandBfor logging andHydrafor configs. Hyperparameters are stored/conf/experimentas.yamlfiles. Depending on your needs, you can specify hyperparameters through the CLI or by changing the.yamlfiles directly.python-mpax.experiment+total_timesteps=1_000_000+num_envs=10We currently support two major environments:MatrixGamesandCoinGame.For `MatrixGames`, we support the ability to specify your own payoff matrix either through the CLI or the `yaml` files. For example the common Iterated Prisoners Dilemma is:
```bash 
python -m pax.experiment +experiment/ipd=ppo ++payoff="[[-2,-2], [0,-3], [-3,0], [-1,-1]]" ++wandb.group="testing"Experimentspython-mpax.experiment+experiment/ipd=yaml++wandb.group="testing"We store previous experiments as parity tests. We useHydrato store these configs and keep track of good hyper-paremters. As a rule for development, we try retain backwards compatability and allow all previous results to be replicated. These can be run easily bypython -m pax.experiment +experiment=NAME. We also provide a  list of our existing experiments and expected resulthere.Loading and SavingAll models trained using Pax by default are saved to theexpfolder.
2a. If you have the model saved locally, specifymodel_path = exp/.... By default, Player 1 will be loaded with the parameters.2b. If you do not have the weights saved locally, specify the wandb runrun_path={wandb-group}{wandb-project}{}andmodel_path = exp/...player 1 will be loaded with the parameters.In order to run evaluation, specifyeval: Trueand evaluation fornum_seedsiterations.CitationIf you use Pax in any of your work, please cite:@misc{pax,
    author = {Khan, Akbir and Willi, Timon and Kwan, Newton, and Samvelyan, Mikayel and Lu, Chris},
    title = {Pax: Multi-Agent Learning in JAX},
    year = {2022},
    publisher = {GitHub},
    journal = {GitHub repository},
    howpublished = {\url{https://github.com/akbir/pax}},
}