Package: minimax-lib
Efficient baselines for autocurricula in JAXContentsWhyminimax?Hardware-accelerated baselinesInstallQuick startDive deeperTrainingLoggingCheckpointingEvaluatingEnvironmentsSupported environmentsAdding environmentsAgentsRoadmapLicenseCitation🐢 Whyminimax?Unsupervised Environment Design (UED) is a promising approach to generating autocurricula for training robust deep reinforcement learning (RL) agents. However, existing implementations of common baselines require excessive amounts of compute. In some cases, experiments can require more than a week to complete using V100 GPUs.This long turn-around slows the rate of research progress in autocuriculum methods.minimaxprovides fast,JAX-basedimplementations of key UED baselines, which are based on the concept ofminimaxregret. By making use of fully-tensorized environment implementations,minimaxbaselines are are fully-jittable and thus take total advantage of the hardware acceleration offered by JAX. In timing studies done on V100 GPUs and Xeon E5-2698 v4 CPUs, we findminimaxbaselines can runover 100x faster than the previous reference implementationson the same single-GPU hardware, like those infacebookresearch/dcd.All autocurriculum algorithms implemented inminimaxalso support multi-device training, which can be activated through asingle command line flag. Using multiple devices for training can lead to further speed ups and allows scaling these autocurriculum methods to much larger batch sizes.🐇 Hardware-accelerated baselinesminimaxincludes JAX-based implementations ofDomain Randomization (DR)Minimax adversaryPAIREDPopulation PAIREDPrioritized Level Replay (PLR)Robust Prioritized Level Replay (PLR$^{\perp}$)ACCELAdditionally,minimaxincludes two new variants of PLR and ACCEL that further reduce wall time by better leveraging the massive degree of environment parallelism enabled by JAX:Parallel PLR (PLR$^{||}$)Parallel ACCEL (ACCEL$^{||}$)In brief, these two new algorithms collect rollouts for new level evaluation, level replay, and, in the case of Parallel ACCEL, mutation evaluation, all in parallel (i.e. rather than sequentially, as done by Robust PLR and ACCEL). As a simple example for why this parallelization improves wall time, consider how Robust PLR with replay probability of0.5would require approximately 2x as many rollouts in order to reach the same number of RL updates as a method like DR, because updates are only performed on rollouts based on level replay. Parallelizing level replay rollouts alongside new level evaluation rollouts by using 2x the environment parallelism reduces the total number of parallel rollouts to equal the total number of updates desired, thereby matching the 1:1 rollout to update ratio of DR. The diagram below summarizes this difference.minimaxincludes a fully-tensorized implementation of a maze environment that we callAMaze. This environment exactly reproduces the MiniGrid-based mazes used in previous UED studies in terms of dynamics, reward function, observation space, and action space, while running many orders of magnitude faster in wall time, with increasing environment parallelism.🛠️ InstallUse a virtual environment manager likecondaormambato create a new environment for your project:condacreate-nminimax
condaactivateminimaxInstallminimaxvia eitherpip install minimax-liborpip install ued.That's it!⚠️ Note that to enable hardware acceleration on GPU, you will need to make sure to install the latest version ofjax>=0.4.19andjaxlib>=0.4.19that is compatible with your CUDA driver (requires minimum CUDA version of11.8). Seethe official JAX installation guidefor detailed instructions.🏁 Quick startThe easiest way to get started is to play with the Python notebooks in theexamples folderof this repository. We also host Colab versions of these notebooks:DR [IPython,Colab]PAIRED [IPython,Colab]PLR and ACCEL*: [IPython,Colab]*Depending on how the top-level flags are set, this notebook runs PLR, Robust PLR, Parallel PLR, ACCEL, or Parallel ACCEL.minimaxcomes with high-performing hyperparameter configurations for several aglorithms, including domain randomization (DR), PAIRED, PLR, and ACCEL for 60-block mazes. You can train using these settings by first creating the training command for executingminimax.trainusing the convenience scriptminimax.config.make_cmd:python -m minimax.config.make_cmd --config maze/[dr,paired,plr,accel] | pbcopy,followed by pasting and executing the resulting command into the command line.See the docsforminimax.config.make_cmdto learn more about how to use this script to generate training commands from JSON configurations. You can browse the available JSON configurations for various autocurriculum methods in theconfigs folder.Note that when logging and checkpointing are enabled, the mainminimax.trainscript outputs this data aslogs.csvandcheckpoint.pklrespectively in an experiment directory located at<log_dir>/<xpid>, wherelog_dirandxpidare arguments specified in the command. You can then evaluate the checkpoint by usingminimax.evaluate:python-mminimax.evaluate\--seed1\--log_dir<absolutepathlogdirectory>\--xpid_prefix<selectcheckpointswithxpidsmatchingthisprefix>\--env_names<csvstringoftestenvironmentnames>\--n_episodes<numberoftrialspertestenvironment>\--results_path<pathtoresultsfolder>\--results_fname<filenameofoutputresultscsv>🪸 Dive deeperTrainingThe main entry for training isminimax.train. This script configures the training run based on command line arguments. It constructs an instance ofExperimentRunnerto manage the training process on an update-cycle basis: These duties include constructing and delegating updates to an appropriate training runner for the specified autocurriculum algorithm and conducting logging and checkpointing. The training runner used byExperimentRunnerexecutes all autocurriculum-related logic. The system diagram above describes how these pieces fit together, as well as howminimaxmanages various, hierarchical batch dimensions.Currently,minimaxincludes training runners for the following classes of autocurricula:RunnerAlgorithm class--train_runnerDRRunnerDomain randomizationdrPLRRunnerReplay-based curricula, including ACCELplrPAIREDRunnerCurricula via a co-adapting teacher environment design policypairedThe below table summarizes how various autocurriculum methods map to these runners and the key arguments that must be set differently from the default settings in order to switch the runner's behavior to each method.AlgorithmReferenceRunnerKey argsDRTobin et al, 2019DRRunner–Minimax adversaryDennis et al, 2020PAIREDRunnerued_score='neg_return'PAIREDDennis et al, 2020PAIREDRunner–Population PAIREDDennis et al, 2020PAIREDRunnern_students >= 2,ued_score='population_regret'PLRJiang et al, 2021PLRRunnerplr_use_robust_plr=FalseRobust PLRJiang et al, 2021aPLRRunner–ACCELParker-Holder et al, 2022PLRRunnerplr_mutation_fn != None,plr_n_mutations > 0Parallel PLRJiang et al, 2023PLRRunnerplr_use_parallel_eval=TrueParallel ACCELJiang et al, 2023PLRRunnerplr_use_parallel_eval=True,plr_mutation_fn != None,plr_n_mutations > 0See the docsonminimax.trainfor a comprehensive guide on how to configure command line arguments for running various autocurricula methods viaminimax.train.LoggingBy default,minimax.traingenerates a folder in the directory specified by the--log_dirargument, named according to--xpid. This folder contains the main training logs,logs.csv, and periodic screenshots of generated levels in the directoryscreenshots. Each screenshot uses the naming conventionupdate_<number of PPO updates>.png. When ACCEL is turned on, the screenshot naming convention also includes information about whether the level was replayed via PLR and the mutation generation number for the level, i.e. how many mutation cycles led to this level.CheckpointingLatest checkpoint:The latest model checkpoint is saved ascheckpoint.pkl. The model is checkpointed every--checkpoint_intervalnumber of updates, where each update corresponds to a full rollout and update cycle for each participating agent. For the same number of environment interaction steps, methods may differ in the number of gradient updates performed by participating agents, so checkpointing based on number of update cycles controls for this potential discrepency. For example, methods based on Robust PLR, like ACCEL, do not perform student gradient updates every rollout cycle.Archived checkpoints:Separate archived model checkpoints can be saved at specific intervals by specifying a positive value for the argument--archive_interval. For example, setting--archive_interval=1000will result in saving model checkpoints for every 1000 updates, namedcheckpoint_1000.tar,checkpoint_2000.tar, and so on. These archived models are saved in addition tocheckpoint.pkl, which always stores the latest checkpoint, based on--checkpoint_interval.EvaluatingOnce training completes, you can evaluate the resultingcheckpoint.pklon test environments usingminimax.eval. This script can evaluate an individual checkpoint or group of checkpoints created via training runs with a shared experiment ID prefix (--xpidvalue), e.g. each corresponding to different training seeds of the same experiment configuration. Each checkpoint is evaluated over--n_episodesepisodes for each of the test environments, specified via a csv string of test environment names passed in via--env_names. The evaluation results can be optionally written to a csv file in--results_path, if a--results_fnameis provided.See the docsonminimax.evaluationfor a comprehensive guide on how to configure command line arguments forminimax.evaluate.Multi-device trainingAll autocurriculum algorithms inminimaxsupport multi-device training viashmapacross the environment batch dimension (see the system diagram above). In order to shard rollouts and gradient updates along the environment batch dimension acrossNdevices, simply passminimax.trainthe additional argument--n_devices=N. By default,n_devices=1.🏝️ EnvironmentsSupported environmentsminimaxcurrently includesAMaze, a fully-tensorized implementation of the partially-observable maze navigation environments featured in previous UED studies (see exampleAMazeenvironments in the figure above). Theminimaximplementation of the maze environment fully replicates the original MiniGrid-based dynamics, reward functions, observation space, action space. See the environment docs fo more details.We look forward to working with the greater RL community in continually expanding the set of environments integrated withminimax.Adding environmentsIn order to integrate intominimax's fully-jittable training logic, environments should be implemented in a tensorized fashion via JAX.  All environments must implement theEnvironmentinterface. At a high level,Environmentsubclasses should implement reset and step logic assuming a single environment instance (no environment parallelism). Parallelism is automatically achieved via the training runner logic included withminimax(See thepaperand system diagram above for a quick overview of how this is performed).A key design decision ofminimaxis to separate environment parameters into two groups:Static parametersare fixed throughout training. These parameters are frozen hyperparameters defining some unchanging aspect of the underlying environment distribution, e.g. the width, height, or maximum number of walls of maze environments considered during training. These static parameters are encapsulated in anEnvParamsdataclass.Free parameterscan change per environment instance (e.g. across each instance in a parallel rollout batch). These parameters might correspond to aspects like the specific wall map defining the maze layout or the starting position of the agent. Free parameters are simply treated as part of the fully-traceableEnvState, taking the form of an arbitrary pytree.All environments supporting theEnvironmentinterface will interoperate withDRRunnerandPLRRunner(though for ACCEL mode, wheremutation_fn != None, a mutation operatormust additionally be defined).Environment design with a co-adapting teacher policyIn PAIRED-style autocurricula, a teacher policy generates environment instances in order to maximize some curriculum objective, e.g. relative regret. The teacher's decision-making process corresponds to its own MDP.To support such autocurricula,minimaxfollows the pattern of implementing separateEnvironmentsubclasses for each of student and teacher MDPs. A convenience class calledUEDEnvironmentis then initialized with instances of the student and teacher MDPs. TheUEDEnvironmentinstance exposes a unified interface for resetting and stepping the teacher and student, which is then used in the training runner. For example, stepping theUEDEnvironmentinstance for the teacher (via thestep_teachermethod) produces an environment instance, which can then be used with thereset_studentmethod to reset the state of the UEDEnvironment object to that particular environment instance. Subsequent calls of thestep_studentmethod then operate within this environment instance. Following this pattern, integration of a new environment withPAIREDRunnerrequires implementing the correspondingEnvironmentsubclass for the teacher MDP (the decision process ). Seeminimax/envs/maze/maze_ued.pyfor an example based on the maze environment.Environment operatorsCustom environment operators can also be defined inminimax.Comparatorstake two environment instances, as represented by theirEnvStatepytrees and returnTrueiff the two instances are deemed equal. If a comparator is registered for an environment, training runners can use the comparator to enforce uniqueness of environment instances for many purposes, e.g. making sure the members of the PLR buffer are all unique.Mutatorstake an environment instance, as represented by anEnvStatepytree, and apply some modification to the instance, returning the modified (or "mutated") instance. Mutators are used by ACCEL to mutate environment instances in the PLR buffer. New environments seeking integration with the ACCEL mode of thePLRRunnershould implement and register a default mutation operator.RegistrationEach newEnvironmentsubclass should be registered with theenvsmodule:Student environmentsshould be registered usingenvs.registration.register. See src/minimax/maze/maze.py for an example.Teacher environmentsshould be registered usingenvs.registration.register_ued. Seeenvs/maze/maze_ued.pyfor an example.Mutatorsshould be registered usingenvs.registration.register_mutator. Seeenvs/maze/maze_mutators.pyfor an example.Comparatorsshould be registered usingenvs.registration.register_comparator. Seeenvs/maze/maze_comparators.pyfor an example.🤖 AgentsInminimaxagentscorrespond to a particular data-seeking learning algorithm, e.g. PPO. Amodelcorresponds to a module that implements the policy (or value function) used by the agent. Any agent that follows theAgentinterface should be usable in anyminimaxcompatible environment.Model forward passes are assumed to return a tuple of(value_prediction, policy_logits, carry).RegistrationCustom model classes should be registered for a particular environment for which they are designed. Seemodels/maze/gridworld_models.pyfor an example. After registration, the model can be easily retrieved and  viamodels.make(env_name, model_name, **model_kwargs).🚀 RoadmapMany exciting features are planned for future releases ofminimax. Features planned for near-term release include:Add support forJaxMARL(multi-agent RL environments) via an IPPO mode inDRRunnerandPLRRunner.ExtendParsnipwith methods for composing argument specs across multiple files to reduce the size ofarguments.pycurrently used fortrain.py.Add support forJumanji(combinatorial optimization environments), via an appropriate decorator class.You can suggest new features or ways to improve current functionality by creating an issue in this repository.🪪 Licenseminimaxis licensed underApache 2.0.📜 CitationFor attribution in academic contexts, please cite this work as@article{jiang2023minimax,
    title={minimax: Efficient Baselines for Autocurricula in JAX},
    author={Jiang, Minqi and Dennis, Michael and Grefenstette, Edward and Rocktäschel, Tim},
    booktitle={Agent Learning in Open-Endedness Workshop at NeurIPS},
    year={2023}
}