Package: xpag
xpag("exploringagents") is a modular reinforcement learning library with JAX agents, currently in beta version.InstallOption 1: conda (preferred option)This option is preferred because it relies mainly on conda-forge packages (which among other things simplifies the installation of JAX).git clone https://github.com/perrin-isir/xpag.git
cd xpag
conda update condaInstall micromamba if you don't already have it (you can also simply use conda, by replacing belowmicromamba create,micromamba updateandmicromamba activaterespectively byconda env create,conda env updateandconda activate, but this will lead to a significantly slower installation):conda install -c conda-forge micromambaChoose an environment name, for instancexpagenv.The following command creates thexpagenvenvironment with the requirements listed inenvironment.yaml:micromamba create --name xpagenv --file environment.yamlIf you prefer to update an existing environment (existing_env):micromamba update --name existing_env --file environment.yamlThen, activate thexpagenvenvironment:micromamba activate xpagenvFinally, install thexpaglibrary in the activated environment:pip install -e .Option 2: pipFor the pip install, you need to properly install JAX yourself. Otherwise, if JAX is installed automatically as a pip dependency ofxpag, it will probably not work as desired (e.g. it will not be GPU-compatible). So you should install it beforehand, following these guidelines:https://github.com/google/jax#installationThen, installxpagwith:pip install xpagJAXTo verify that the JAX installation went well, check the backend used by JAX with the following command:python -c "import jax; print(jax.lib.xla_bridge.get_backend().platform)"It will print "cpu", "gpu" or "tpu" depending on the platform JAX is using.TutorialsThe following libraries, not required byxpag, are required for thetutorials:MuJoCo (pip install mujoco): seehttps://github.com/deepmind/mujocoimageio (pip install imageio): seehttps://github.com/imageio/imageioTutorialsThexpag-tutorialsrepository contains a list of tutorials (colab notebooks) forxpag:https://github.com/perrin-isir/xpag-tutorialsShort documentationxpag: a platform for RL, goal-conditioned RL, and more.xpagallows standard reinforcement learning, but it has been designed with
goal-conditioned reinforcement learning (GCRL) in mind (check out thetrain_gmazes.ipynbtutorial for a simple example of GCRL).In GCRL, agents have a goal, which is part of the input they take, and the reward mainly depends on
the degree of achievement of that goal. Beyond the usual modules in
RL platforms (environment, agent, buffer/sampler),xpagintroduces a
module called "setter" which, among other things, can help to set and manage
goals (for example modifying the goal several times in a single episode).
Although the setter is largely similar to an environment wrapper, it
is separated from the environment because in some cases it should be considered as
an independent entity (e.g. a teacher), or as a part of the agent itself.xpagrelies on a single reinforcement learning loop (thelearn()function inxpag/tools/learn.py)
in which the environment, the agent, the buffer and the setter interact (see below).
Thelearn()function  has the following first 3 arguments (returned bygym_vec_env()andbrax_vec_env()):env: the training environment, which runs 1 or more rollouts in parallel.eval_env: the evaluation environment, identical toenvexcept that it runs
a single rollout.env_info: a dictionary containing information about the environment:env_info["env_type"]: the type of environment; for the momentxpagdifferentiates 3 types of environments: "Brax" environments, "Mujoco" environments, and
"Gym" environments. This information is used to adapt the way episodes are saved and replayed.env_info["name"]: the name of the environment.env_info["is_goalenv"]: whether the environment is a goal-based environment or
not.env_info["num_envs"]: the number of parallel rollouts inenvenv_info["max_episode_steps"]: the maximum number of steps in episodes (xpagdoes not allow potentially infinite episodes).env_info["action_space"]: the action space (of typegym.spaces.Space) that takes into account parallel rollouts. It can be useful to sample random actions.env_info["single_action_space"]: the action space (of typegym.spaces.Space) for single rollouts.learn()also takes in input the agent, the buffer and the setter and various parameters. Detailed information about the arguments oflearn()can be
found in the code documentation (checkxpag/tools/learn.py).The components that interact during learning are:the environment (env)Inxpag, environments must allow parallel rollouts, andxpagkeeps the same API even in the case of a single rollout,
i.e. when the number of "parallel environments" is 1. Basically, all environments are
"vector environments".env.reset(seed: Optional[Union[int, List[int]]], options: Optional[dict])->observation: Union[np.array, jax.numpy.array], info: dictFollowing the gym Vector API
(seehttps://www.gymlibrary.dev/api/vector/#vectorenv), environments have
areset()function that returns anobservation(which is actually a batch of observations for all the
parallel rollouts) and an optional dictionaryinfo(seehttps://www.gymlibrary.dev/api/vector/#reset).We expectobservationto be a numpy array, or a jax.numpy array, and its first dimension
selects between parallel rollouts, which means thatobservation[i]is the observation in
the i-th rollout. In the case of a single rollout,observation[0]is the observation
in this rollout.env.step(action: Union[np.array, jax.numpy.array])->observation, reward, terminated, truncated, infoAgain, following the gym Vector API, environments have astep()function that takes
in input an action (which is actually a batch of actions, one per rollout) and returns:observation,reward,terminated,truncated,info(cf.https://www.gymlibrary.dev/api/vector/#step).
There are slight differences with the gym Vector API. First, inxpagthis API also covers the case
of a single rollout. Second,xpagassumes thatreward,terminatedandtruncatedhave shape(n, 1), not(n,)(where n is the number of parallel rollouts). More broadly, whether they are due to a single rollout or to
unidimensional elements, single-dimensional entries are not squeezed inxpag.
Third, inxpag,infois a dictionary, not a tuple of dictionaries
(however its entries may be tuples).env.reset_done(done, seed: Optional[Union[int, List[int]]], options: Optional[dict])->observation, infoThe most significant difference with the gym Vector API is thatxpagrequires areset_done()function which takes adonearray of Booleans in input and performs a reset for
the i-th rollout if and only ifdone[i]is evaluated to True. Besidesdone, the arguments ofreset_done()are the same as the ones ofreset():seedandoptions, and its outputs are also the same:observation,info.
For rollouts that are not reset, the returned observation is the same as the observation returned by the laststep().reset()must be called once for the initial reset, and afterwards onlyreset_done()should be used. Auto-resets (automatic resets after terminal transitions) are not allowed inxpag.
The main reason to preferreset_done()to auto-resets
is that with auto-resets, terminal transitions must be special and contain additional
information. Withreset_done(), this is no longer necessary. Furthermore,
by modifying thedonearray returned by a step of the environment, it becomes possible
to easily force the termination of an episode, or to force an episode to continue despite
reaching a terminal transition (but this must be done with caution).gym_vec_env(env_name: str, num_envs: int, wrap_function: Callable = None)->env, eval_env, env_info: dictbrax_vec_env(env_name: str, num_envs: int, wrap_function: Callable = None, *, force_cpu_backend : bool = False)->env, eval_env, env_info: dictThegym_vec_env()andbrax_vec_env()functions (seetutorials)
call wrappers that automatically add thereset_done()function to Gym and Brax
environments, and make the wrapped environments fit thexpagAPI.Goal-based environments:Goal-based environments (for GCRL) must have a similar interface to the one defined in
theGym-Roboticslibrary
(seeGoalEnvincore.py), with minor differences.
Their observation spaces are of typegym.spaces.Dict, with the following keys
in theobservationdictionaries:"observation","achieved_goal", and"desired_goal".
Goal-based environments must also have in attribute acompute_reward()function that computes rewards.
Inxpag, the inputs ofcompute_reward()can be different from the ones considered in
the originalGoalEnvclass. For example, in theGoalEnvWrapperclass,
which can be used to turn standard environments into goal-based environments, the
arguments ofcompute_reward()are assumed to beachieved_goal(the goal achievedafterstep()),desired_goal(the desired goalbeforestep()),action,observation(the observationafterstep()),reward(the reward of the base environment),terminated,truncatedandinfo(the outputs of thestep()function). In the version ofHER(cf.https://arxiv.org/pdf/1707.01495.pdf) inxpag,
it is assumed thatcompute_reward()depends only onachieved_goal,desired_goal,actionandobservation.In goal-based environments, the multiple observations from parallel rollouts are concatenated as in the gym functionconcatenate()(cf.https://github.com/openai/gym/blob/master/gym/vector/utils/numpy_utils.py),
which means that the batched observations are always single dictionaries in which the
entries"observation","achieved_goal"and"desired_goal"are arrays of observations,
achieved goals and desired goals.infoxpagassumes that, in goal-based environments, theinfodictionary returned bystep()always containsinfo["is_success"], an array of Booleans (one per rollout)
that areTrueif the corresponding transition is a successfull achievement of the
desired goal, andFalseotherwise (remark:this does not need to coincide
with episode termination).the agent (agent)xpagonly considers off-policy agents. (TODO)the buffer (buffer)TODOthe sampler (sampler)TODOthe setter (setter)TODOThe figure below summarizes the RL loop and the interactions between the components:
(TODO)AcknowledgementsMaintainer and main contributor:Nicolas Perrin-Gilbert (CNRS, ISIR)Other people who contributed toxpag:Olivier Serris (ISIR)Alexandre Chenu (ISIR)Stéphane Caron (Inria)Fabian Schramm (Inria)There is an interface to agents from theRLJAXlibrary (seerljax_interface.py). This provideshaikuversions ofDDPG,TD3,TQC,SACand SAC withDisCor.Theflaxversion of theSAC agentis based on the implementation of SAC inJAXRL, and some elements of the flax version of theTQC agentcome from the implementation of TQC inRLJAX.Citing the projectTo cite this repository in publications:@misc{xpag,author={Perrin-Gilbert, Nicolas},title={xpag: a modular reinforcement learning library with JAX agents},year={2022},url={https://github.com/perrin-isir/xpag}}