discount_aggregator |
Callable |
a function to aggregate all agents discounts into a single
scalar value, e.g. max. |
- <function amax at 0x7f157acdeee0> |
+ <function amax at 0x7fe959e73ee0> |
diff --git a/search/search_index.json b/search/search_index.json
index 03db6f660..e74ec0bef 100644
--- a/search/search_index.json
+++ b/search/search_index.json
@@ -1 +1 @@
-{"config":{"indexing":"full","lang":["en"],"min_search_length":3,"prebuild_index":false,"separator":"[\\s\\-]+"},"docs":[{"location":"","text":"Environments | Installation | Quickstart | Training | Citation | Docs Welcome to the Jungle! \ud83c\udf34 # Jumanji is a diverse suite of scalable reinforcement learning environments written in JAX. Jumanji is helping pioneer a new wave of hardware-accelerated research and development in the field of RL. Jumanji's high-speed environments enable faster iteration and large-scale experimentation while simultaneously reducing complexity. Originating in the Research Team at InstaDeep , Jumanji is now developed jointly with the open-source community. To join us in these efforts, reach out, raise issues and read our contribution guidelines or just star \ud83c\udf1f to stay up to date with the latest developments! Goals \ud83d\ude80 # Provide a simple, well-tested API for JAX-based environments. Make research in RL more accessible. Facilitate the research on RL for problems in the industry and help close the gap between research and industrial applications. Provide environments whose difficulty can be scaled to be arbitrarily hard. Overview \ud83e\udd9c # \ud83e\udd51 Environment API : core abstractions for JAX-based environments. \ud83d\udd79\ufe0f Environment Suite : a collection of RL environments ranging from simple games to NP-hard combinatorial problems. \ud83c\udf6c Wrappers : easily connect to your favourite RL frameworks and libraries such as Acme , Stable Baselines3 , RLlib , OpenAI Gym and DeepMind-Env through our dm_env and gym wrappers. \ud83c\udf93 Examples : guides to facilitate Jumanji's adoption and highlight the added value of JAX-based environments. \ud83c\udfce\ufe0f Training: example agents that can be used as inspiration for the agents one may implement in their research. Environments \ud83c\udf0d Jumanji provides a diverse range of environments ranging from simple games to NP-hard combinatorial problems. Environment Category Registered Version(s) Source Description \ud83d\udd22 Game2048 Logic Game2048-v1 code doc \ud83c\udfa8 GraphColoring Logic GraphColoring-v0 code doc \ud83d\udca3 Minesweeper Logic Minesweeper-v0 code doc \ud83c\udfb2 RubiksCube Logic RubiksCube-v0 RubiksCube-partly-scrambled-v0 code doc \u270f\ufe0f Sudoku Logic Sudoku-v0 Sudoku-very-easy-v0 code doc \ud83d\udce6 BinPack (3D BinPacking Problem) Packing BinPack-v2 code doc \ud83c\udfed JobShop (Job Shop Scheduling Problem) Packing JobShop-v0 code doc \ud83c\udf92 Knapsack Packing Knapsack-v1 code doc \u2592 Tetris Packing Tetris-v0 code doc \ud83e\uddf9 Cleaner Routing Cleaner-v0 code doc Connector Routing Connector-v2 code doc \ud83d\ude9a CVRP (Capacitated Vehicle Routing Problem) Routing CVRP-v1 code doc \ud83d\ude9a MultiCVRP (Multi-Agent Capacitated Vehicle Routing Problem) Routing MultiCVRP-v0 code doc Maze Routing Maze-v0 code doc RobotWarehouse Routing RobotWarehouse-v0 code doc \ud83d\udc0d Snake Routing Snake-v1 code doc \ud83d\udcec TSP (Travelling Salesman Problem) Routing TSP-v1 code doc Multi Minimum Spanning Tree Problem Routing MMST-v0 code doc Installation \ud83c\udfac You can install the latest release of Jumanji from PyPI: 1 pip install jumanji Alternatively, you can install the latest development version directly from GitHub: 1 pip install git+https://github.com/instadeepai/jumanji.git Jumanji has been tested on Python 3.8 and 3.9. Note that because the installation of JAX differs depending on your hardware accelerator, we advise users to explicitly install the correct JAX version (see the official installation guide ). Rendering: Matplotlib is used for rendering all the environments. To visualize the environments you will need a GUI backend. For example, on Linux, you can install Tk via: apt-get install python3-tk , or using conda: conda install tk . Check out Matplotlib backends for a list of backends you can use. Quickstart \u26a1 RL practitioners will find Jumanji's interface familiar as it combines the widely adopted OpenAI Gym and DeepMind Environment interfaces. From OpenAI Gym, we adopted the idea of a registry and the render method, while our TimeStep structure is inspired by DeepMind Environment. Basic Usage \ud83e\uddd1\u200d\ud83d\udcbb # 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 import jax import jumanji # Instantiate a Jumanji environment using the registry env = jumanji . make ( 'Snake-v1' ) # Reset your (jit-able) environment key = jax . random . PRNGKey ( 0 ) state , timestep = jax . jit ( env . reset )( key ) # (Optional) Render the env state env . render ( state ) # Interact with the (jit-able) environment action = env . action_spec () . generate_value () # Action selection (dummy value here) state , timestep = jax . jit ( env . step )( state , action ) # Take a step and observe the next state and time step state represents the internal state of the environment: it contains all the information required to take a step when executing an action. This should not be confused with the observation contained in the timestep , which is the information perceived by the agent. timestep is a dataclass containing step_type , reward , discount , observation and extras . This structure is similar to dm_env.TimeStep except for the extras field that was added to allow users to log environments metrics that are neither part of the agent's observation nor part of the environment's internal state. Advanced Usage \ud83e\uddd1\u200d\ud83d\udd2c # Being written in JAX, Jumanji's environments benefit from many of its features including automatic vectorization/parallelization ( jax.vmap , jax.pmap ) and JIT-compilation ( jax.jit ), which can be composed arbitrarily. We provide an example of a more advanced usage in the advanced usage guide . Registry and Versioning \ud83d\udcd6 # Like OpenAI Gym, Jumanji keeps a strict versioning of its environments for reproducibility reasons. We maintain a registry of standard environments with their configuration. For each environment, a version suffix is appended, e.g. Snake-v1 . When changes are made to environments that might impact learning results, the version number is incremented by one to prevent potential confusion. For a full list of registered versions of each environment, check out the documentation . Training \ud83c\udfce\ufe0f To showcase how to train RL agents on Jumanji environments, we provide a random agent and a vanilla actor-critic (A2C) agent. These agents can be found in jumanji/training/ . Because the environment framework in Jumanji is so flexible, it allows pretty much any problem to be implemented as a Jumanji environment, giving rise to very diverse observations. For this reason, environment-specific networks are required to capture the symmetries of each environment. Alongside the A2C agent implementation, we provide examples of such environment-specific actor-critic networks in jumanji/training/networks . \u26a0\ufe0f The example agents in jumanji/training are only meant to serve as inspiration for how one can implement an agent. Jumanji is first and foremost a library of environments - as such, the agents and networks will not be maintained to a production standard. For more information on how to use the example agents, see the training guide . Contributing \ud83e\udd1d # Contributions are welcome! See our issue tracker for good first issues . Please read our contributing guidelines for details on how to submit pull requests, our Contributor License Agreement, and community guidelines. Citing Jumanji \u270f\ufe0f If you use Jumanji in your work, please cite the library using: 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 @misc{bonnet2023jumanji, title={Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX}, author={ Cl\u00e9ment Bonnet and Daniel Luo and Donal Byrne and Shikha Surana and Vincent Coyette and Paul Duckworth and Laurence I. Midgley and Tristan Kalloniatis and Sasha Abramowitz and Cemlyn N. Waters and Andries P. Smit and Nathan Grinsztajn and Ulrich A. Mbou Sob and Omayma Mahjoub and Elshadai Tegegn and Mohamed A. Mimouni and Raphael Boige and Ruan de Kock and Daniel Furelos-Blanco and Victor Le and Arnu Pretorius and Alexandre Laterre }, year={2023}, eprint={2306.09884}, url={https://arxiv.org/abs/2306.09884}, archivePrefix={arXiv}, primaryClass={cs.LG} } See Also \ud83d\udd0e # Other works have embraced the approach of writing RL environments in JAX. In particular, we suggest users check out the following sister repositories: \ud83e\udd16 Qdax is a library to accelerate Quality-Diversity and neuro-evolution algorithms through hardware accelerators and parallelization. \ud83c\udf33 Evojax provides tools to enable neuroevolution algorithms to work with neural networks running across multiple TPU/GPUs. \ud83e\uddbe Brax is a differentiable physics engine that simulates environments made up of rigid bodies, joints, and actuators. \ud83c\udfcb\ufe0f\u200d Gymnax implements classic environments including classic control, bsuite, MinAtar and a collection of meta RL tasks. \ud83c\udfb2 Pgx provides classic board game environments like Backgammon, Shogi, and Go. Acknowledgements \ud83d\ude4f # The development of this library was supported with Cloud TPUs from Google's TPU Research Cloud (TRC) \ud83c\udf24.","title":"Home"},{"location":"#welcome-to-the-jungle","text":"Jumanji is a diverse suite of scalable reinforcement learning environments written in JAX. Jumanji is helping pioneer a new wave of hardware-accelerated research and development in the field of RL. Jumanji's high-speed environments enable faster iteration and large-scale experimentation while simultaneously reducing complexity. Originating in the Research Team at InstaDeep , Jumanji is now developed jointly with the open-source community. To join us in these efforts, reach out, raise issues and read our contribution guidelines or just star \ud83c\udf1f to stay up to date with the latest developments!","title":"Welcome to the Jungle! \ud83c\udf34"},{"location":"#goals","text":"Provide a simple, well-tested API for JAX-based environments. Make research in RL more accessible. Facilitate the research on RL for problems in the industry and help close the gap between research and industrial applications. Provide environments whose difficulty can be scaled to be arbitrarily hard.","title":"Goals \ud83d\ude80"},{"location":"#overview","text":"\ud83e\udd51 Environment API : core abstractions for JAX-based environments. \ud83d\udd79\ufe0f Environment Suite : a collection of RL environments ranging from simple games to NP-hard combinatorial problems. \ud83c\udf6c Wrappers : easily connect to your favourite RL frameworks and libraries such as Acme , Stable Baselines3 , RLlib , OpenAI Gym and DeepMind-Env through our dm_env and gym wrappers. \ud83c\udf93 Examples : guides to facilitate Jumanji's adoption and highlight the added value of JAX-based environments. \ud83c\udfce\ufe0f Training: example agents that can be used as inspiration for the agents one may implement in their research.","title":"Overview \ud83e\udd9c"},{"location":"#basic-usage","text":"1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 import jax import jumanji # Instantiate a Jumanji environment using the registry env = jumanji . make ( 'Snake-v1' ) # Reset your (jit-able) environment key = jax . random . PRNGKey ( 0 ) state , timestep = jax . jit ( env . reset )( key ) # (Optional) Render the env state env . render ( state ) # Interact with the (jit-able) environment action = env . action_spec () . generate_value () # Action selection (dummy value here) state , timestep = jax . jit ( env . step )( state , action ) # Take a step and observe the next state and time step state represents the internal state of the environment: it contains all the information required to take a step when executing an action. This should not be confused with the observation contained in the timestep , which is the information perceived by the agent. timestep is a dataclass containing step_type , reward , discount , observation and extras . This structure is similar to dm_env.TimeStep except for the extras field that was added to allow users to log environments metrics that are neither part of the agent's observation nor part of the environment's internal state.","title":"Basic Usage \ud83e\uddd1\u200d\ud83d\udcbb"},{"location":"#advanced-usage","text":"Being written in JAX, Jumanji's environments benefit from many of its features including automatic vectorization/parallelization ( jax.vmap , jax.pmap ) and JIT-compilation ( jax.jit ), which can be composed arbitrarily. We provide an example of a more advanced usage in the advanced usage guide .","title":"Advanced Usage \ud83e\uddd1\u200d\ud83d\udd2c"},{"location":"#registry-and-versioning","text":"Like OpenAI Gym, Jumanji keeps a strict versioning of its environments for reproducibility reasons. We maintain a registry of standard environments with their configuration. For each environment, a version suffix is appended, e.g. Snake-v1 . When changes are made to environments that might impact learning results, the version number is incremented by one to prevent potential confusion. For a full list of registered versions of each environment, check out the documentation .","title":"Registry and Versioning \ud83d\udcd6"},{"location":"#contributing","text":"Contributions are welcome! See our issue tracker for good first issues . Please read our contributing guidelines for details on how to submit pull requests, our Contributor License Agreement, and community guidelines.","title":"Contributing \ud83e\udd1d"},{"location":"#see-also","text":"Other works have embraced the approach of writing RL environments in JAX. In particular, we suggest users check out the following sister repositories: \ud83e\udd16 Qdax is a library to accelerate Quality-Diversity and neuro-evolution algorithms through hardware accelerators and parallelization. \ud83c\udf33 Evojax provides tools to enable neuroevolution algorithms to work with neural networks running across multiple TPU/GPUs. \ud83e\uddbe Brax is a differentiable physics engine that simulates environments made up of rigid bodies, joints, and actuators. \ud83c\udfcb\ufe0f\u200d Gymnax implements classic environments including classic control, bsuite, MinAtar and a collection of meta RL tasks. \ud83c\udfb2 Pgx provides classic board game environments like Backgammon, Shogi, and Go.","title":"See Also \ud83d\udd0e"},{"location":"#acknowledgements","text":"The development of this library was supported with Cloud TPUs from Google's TPU Research Cloud (TRC) \ud83c\udf24.","title":"Acknowledgements \ud83d\ude4f"},{"location":"api/env/","text":"Environment ( ABC , Generic ) # Environment written in Jax that differs from the gym API to make the step and reset functions jittable. The state contains all the dynamics and data needed to step the environment, no computation stored in attributes of self. The API is inspired by brax . unwrapped : Environment property readonly # reset ( self , key : PRNGKeyArray ) -> Tuple [ ~ State , jumanji . types . TimeStep ] # Resets the environment to an initial state. Parameters: Name Type Description Default key PRNGKeyArray random key used to reset the environment. required Returns: Type Description state State object corresponding to the new state of the environment, timestep: TimeStep object corresponding the first timestep returned by the environment, step ( self , state : ~ State , action : Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ]) -> Tuple [ ~ State , jumanji . types . TimeStep ] # Run one timestep of the environment's dynamics. Parameters: Name Type Description Default state ~State State object containing the dynamics of the environment. required action Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] Array containing the action to take. required Returns: Type Description state State object corresponding to the next state of the environment, timestep: TimeStep object corresponding the timestep returned by the environment, observation_spec ( self ) -> Spec # Returns the observation spec. Returns: Type Description observation_spec a NestedSpec tree of spec. action_spec ( self ) -> Spec # Returns the action spec. Returns: Type Description action_spec a NestedSpec tree of spec. reward_spec ( self ) -> Array # Describes the reward returned by the environment. By default, this is assumed to be a single float. Returns: Type Description reward_spec a specs.Array spec. discount_spec ( self ) -> BoundedArray # Describes the discount returned by the environment. By default, this is assumed to be a single float between 0 and 1. Returns: Type Description discount_spec a specs.BoundedArray spec. render ( self , state : ~ State ) -> Any # Render frames of the environment for a given state. Parameters: Name Type Description Default state ~State State object containing the current dynamics of the environment. required close ( self ) -> None # Perform any necessary cleanup. __enter__ ( self ) -> Environment special # __exit__ ( self , * args : Any ) -> None special # Calls :meth: close() .","title":"Base"},{"location":"api/env/#jumanji.env.Environment","text":"Environment written in Jax that differs from the gym API to make the step and reset functions jittable. The state contains all the dynamics and data needed to step the environment, no computation stored in attributes of self. The API is inspired by brax .","title":"Environment"},{"location":"api/env/#jumanji.env.Environment.unwrapped","text":"","title":"unwrapped"},{"location":"api/env/#jumanji.env.Environment.reset","text":"Resets the environment to an initial state. Parameters: Name Type Description Default key PRNGKeyArray random key used to reset the environment. required Returns: Type Description state State object corresponding to the new state of the environment, timestep: TimeStep object corresponding the first timestep returned by the environment,","title":"reset()"},{"location":"api/env/#jumanji.env.Environment.step","text":"Run one timestep of the environment's dynamics. Parameters: Name Type Description Default state ~State State object containing the dynamics of the environment. required action Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] Array containing the action to take. required Returns: Type Description state State object corresponding to the next state of the environment, timestep: TimeStep object corresponding the timestep returned by the environment,","title":"step()"},{"location":"api/env/#jumanji.env.Environment.observation_spec","text":"Returns the observation spec. Returns: Type Description observation_spec a NestedSpec tree of spec.","title":"observation_spec()"},{"location":"api/env/#jumanji.env.Environment.action_spec","text":"Returns the action spec. Returns: Type Description action_spec a NestedSpec tree of spec.","title":"action_spec()"},{"location":"api/env/#jumanji.env.Environment.reward_spec","text":"Describes the reward returned by the environment. By default, this is assumed to be a single float. Returns: Type Description reward_spec a specs.Array spec.","title":"reward_spec()"},{"location":"api/env/#jumanji.env.Environment.discount_spec","text":"Describes the discount returned by the environment. By default, this is assumed to be a single float between 0 and 1. Returns: Type Description discount_spec a specs.BoundedArray spec.","title":"discount_spec()"},{"location":"api/env/#jumanji.env.Environment.render","text":"Render frames of the environment for a given state. Parameters: Name Type Description Default state ~State State object containing the current dynamics of the environment. required","title":"render()"},{"location":"api/env/#jumanji.env.Environment.close","text":"Perform any necessary cleanup.","title":"close()"},{"location":"api/env/#jumanji.env.Environment.__enter__","text":"","title":"__enter__()"},{"location":"api/env/#jumanji.env.Environment.__exit__","text":"Calls :meth: close() .","title":"__exit__()"},{"location":"api/types/","text":"types # StepType ( int8 ) # Defines the status of a TimeStep within a sequence. First: 0 Mid: 1 Last: 2 TimeStep ( Generic , Mapping ) dataclass # Copied from dm_env.TimeStep with the goal of making it a Jax Type. The original dm_env.TimeStep is not a Jax type because inheriting a namedtuple is not treated as a valid Jax type (https://github.com/google/jax/issues/806). A TimeStep contains the data emitted by an environment at each step of interaction. A TimeStep holds a step_type , an observation (typically a NumPy array or a dict or list of arrays), and an associated reward and discount . The first TimeStep in a sequence will have StepType.FIRST . The final TimeStep will have StepType.LAST . All other TimeStep s in a sequence will have `StepType.MID. Attributes: Name Type Description step_type StepType A StepType enum value. reward Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] A scalar, NumPy array, nested dict, list or tuple of rewards; or None if step_type is StepType.FIRST , i.e. at the start of a sequence. discount Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] A scalar, NumPy array, nested dict, list or tuple of discount values in the range [0, 1] , or None if step_type is StepType.FIRST , i.e. at the start of a sequence. observation ~Observation A NumPy array, or a nested dict, list or tuple of arrays. Scalar values that can be cast to NumPy arrays (e.g. Python floats) are also valid in place of a scalar array. extras Optional[Dict] environment metric(s) or information returned by the environment but not observed by the agent (hence not in the observation). For example, it could be whether an invalid action was taken. In most environments, extras is None. step_type : StepType dataclass-field # reward : Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ] dataclass-field # discount : Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ] dataclass-field # observation : ~ Observation dataclass-field # extras : Optional [ Dict ] dataclass-field # __eq__ ( self , other ) special # __init__ ( self , step_type : StepType , reward : Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ], discount : Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ], observation : ~ Observation , extras : Optional [ Dict ] = None ) -> None special # __repr__ ( self ) special # __getitem__ ( self , x ) special # __len__ ( self ) special # __iter__ ( self ) special # first ( self ) -> Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ] # mid ( self ) -> Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ] # last ( self ) -> Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ] # from_tuple ( args ) # to_tuple ( self ) # replace ( self , ** kwargs ) # __getstate__ ( self ) special # __setstate__ ( self , state ) special # restart ( observation : ~ Observation , extras : Optional [ Dict ] = None , shape : Union [ int , Sequence [ int ]] = ()) -> TimeStep # Returns a TimeStep with step_type set to StepType.FIRST . Parameters: Name Type Description Default observation ~Observation array or tree of arrays. required extras Optional[Dict] environment metric(s) or information returned by the environment but not observed by the agent (hence not in the observation). For example, it could be whether an invalid action was taken. In most environments, extras is None. None shape Union[int, Sequence[int]] optional parameter to specify the shape of the rewards and discounts. Allows multi-agent environment compatibility. Defaults to () for scalar reward and discount. () Returns: Type Description TimeStep TimeStep identified as a reset. transition ( reward : Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ], observation : ~ Observation , discount : Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ] = None , extras : Optional [ Dict ] = None , shape : Union [ int , Sequence [ int ]] = ()) -> TimeStep # Returns a TimeStep with step_type set to StepType.MID . Parameters: Name Type Description Default reward Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] array. required observation ~Observation array or tree of arrays. required discount Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] array. None extras Optional[Dict] environment metric(s) or information returned by the environment but not observed by the agent (hence not in the observation). For example, it could be whether an invalid action was taken. In most environments, extras is None. None shape Union[int, Sequence[int]] optional parameter to specify the shape of the rewards and discounts. Allows multi-agent environment compatibility. Defaults to () for scalar reward and discount. () Returns: Type Description TimeStep TimeStep identified as a transition. termination ( reward : Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ], observation : ~ Observation , extras : Optional [ Dict ] = None , shape : Union [ int , Sequence [ int ]] = ()) -> TimeStep # Returns a TimeStep with step_type set to StepType.LAST . Parameters: Name Type Description Default reward Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] array. required observation ~Observation array or tree of arrays. required extras Optional[Dict] environment metric(s) or information returned by the environment but not observed by the agent (hence not in the observation). For example, it could be whether an invalid action was taken. In most environments, extras is None. None shape Union[int, Sequence[int]] optional parameter to specify the shape of the rewards and discounts. Allows multi-agent environment compatibility. Defaults to () for scalar reward and discount. () Returns: Type Description TimeStep TimeStep identified as the termination of an episode. truncation ( reward : Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ], observation : ~ Observation , discount : Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ] = None , extras : Optional [ Dict ] = None , shape : Union [ int , Sequence [ int ]] = ()) -> TimeStep # Returns a TimeStep with step_type set to StepType.LAST . Parameters: Name Type Description Default reward Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] array. required observation ~Observation array or tree of arrays. required discount Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] array. None extras Optional[Dict] environment metric(s) or information returned by the environment but not observed by the agent (hence not in the observation). For example, it could be whether an invalid action was taken. In most environments, extras is None. None shape Union[int, Sequence[int]] optional parameter to specify the shape of the rewards and discounts. Allows multi-agent environment compatibility. Defaults to () for scalar reward and discount. () Returns: Type Description TimeStep TimeStep identified as the truncation of an episode. get_valid_dtype ( dtype : Union [ numpy . dtype , type ]) -> dtype # Cast a dtype taking into account the user type precision. E.g., if 64 bit is not enabled, jnp.dtype(jnp.float_) is still float64. By passing the given dtype through jnp.empty we get the supported dtype of float32. Parameters: Name Type Description Default dtype Union[numpy.dtype, type] jax numpy dtype or string specifying the array dtype. required Returns: Type Description dtype dtype converted to the correct type precision.","title":"Types"},{"location":"api/types/#jumanji.types","text":"","title":"types"},{"location":"api/types/#jumanji.types.StepType","text":"Defines the status of a TimeStep within a sequence. First: 0 Mid: 1 Last: 2","title":"StepType"},{"location":"api/types/#jumanji.types.TimeStep","text":"Copied from dm_env.TimeStep with the goal of making it a Jax Type. The original dm_env.TimeStep is not a Jax type because inheriting a namedtuple is not treated as a valid Jax type (https://github.com/google/jax/issues/806). A TimeStep contains the data emitted by an environment at each step of interaction. A TimeStep holds a step_type , an observation (typically a NumPy array or a dict or list of arrays), and an associated reward and discount . The first TimeStep in a sequence will have StepType.FIRST . The final TimeStep will have StepType.LAST . All other TimeStep s in a sequence will have `StepType.MID. Attributes: Name Type Description step_type StepType A StepType enum value. reward Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] A scalar, NumPy array, nested dict, list or tuple of rewards; or None if step_type is StepType.FIRST , i.e. at the start of a sequence. discount Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] A scalar, NumPy array, nested dict, list or tuple of discount values in the range [0, 1] , or None if step_type is StepType.FIRST , i.e. at the start of a sequence. observation ~Observation A NumPy array, or a nested dict, list or tuple of arrays. Scalar values that can be cast to NumPy arrays (e.g. Python floats) are also valid in place of a scalar array. extras Optional[Dict] environment metric(s) or information returned by the environment but not observed by the agent (hence not in the observation). For example, it could be whether an invalid action was taken. In most environments, extras is None.","title":"TimeStep"},{"location":"api/types/#jumanji.types.restart","text":"Returns a TimeStep with step_type set to StepType.FIRST . Parameters: Name Type Description Default observation ~Observation array or tree of arrays. required extras Optional[Dict] environment metric(s) or information returned by the environment but not observed by the agent (hence not in the observation). For example, it could be whether an invalid action was taken. In most environments, extras is None. None shape Union[int, Sequence[int]] optional parameter to specify the shape of the rewards and discounts. Allows multi-agent environment compatibility. Defaults to () for scalar reward and discount. () Returns: Type Description TimeStep TimeStep identified as a reset.","title":"restart()"},{"location":"api/types/#jumanji.types.transition","text":"Returns a TimeStep with step_type set to StepType.MID . Parameters: Name Type Description Default reward Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] array. required observation ~Observation array or tree of arrays. required discount Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] array. None extras Optional[Dict] environment metric(s) or information returned by the environment but not observed by the agent (hence not in the observation). For example, it could be whether an invalid action was taken. In most environments, extras is None. None shape Union[int, Sequence[int]] optional parameter to specify the shape of the rewards and discounts. Allows multi-agent environment compatibility. Defaults to () for scalar reward and discount. () Returns: Type Description TimeStep TimeStep identified as a transition.","title":"transition()"},{"location":"api/types/#jumanji.types.termination","text":"Returns a TimeStep with step_type set to StepType.LAST . Parameters: Name Type Description Default reward Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] array. required observation ~Observation array or tree of arrays. required extras Optional[Dict] environment metric(s) or information returned by the environment but not observed by the agent (hence not in the observation). For example, it could be whether an invalid action was taken. In most environments, extras is None. None shape Union[int, Sequence[int]] optional parameter to specify the shape of the rewards and discounts. Allows multi-agent environment compatibility. Defaults to () for scalar reward and discount. () Returns: Type Description TimeStep TimeStep identified as the termination of an episode.","title":"termination()"},{"location":"api/types/#jumanji.types.truncation","text":"Returns a TimeStep with step_type set to StepType.LAST . Parameters: Name Type Description Default reward Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] array. required observation ~Observation array or tree of arrays. required discount Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] array. None extras Optional[Dict] environment metric(s) or information returned by the environment but not observed by the agent (hence not in the observation). For example, it could be whether an invalid action was taken. In most environments, extras is None. None shape Union[int, Sequence[int]] optional parameter to specify the shape of the rewards and discounts. Allows multi-agent environment compatibility. Defaults to () for scalar reward and discount. () Returns: Type Description TimeStep TimeStep identified as the truncation of an episode.","title":"truncation()"},{"location":"api/types/#jumanji.types.get_valid_dtype","text":"Cast a dtype taking into account the user type precision. E.g., if 64 bit is not enabled, jnp.dtype(jnp.float_) is still float64. By passing the given dtype through jnp.empty we get the supported dtype of float32. Parameters: Name Type Description Default dtype Union[numpy.dtype, type] jax numpy dtype or string specifying the array dtype. required Returns: Type Description dtype dtype converted to the correct type precision.","title":"get_valid_dtype()"},{"location":"api/wrappers/","text":"wrappers # Wrapper ( Environment , Generic ) # Wraps the environment to allow modular transformations. Source: https://github.com/google/brax/blob/main/brax/envs/env.py#L72 unwrapped : Environment property readonly # Returns the wrapped env. __init__ ( self , env : Environment ) special # reset ( self , key : PRNGKeyArray ) -> Tuple [ ~ State , jumanji . types . TimeStep ] # Resets the environment to an initial state. Parameters: Name Type Description Default key PRNGKeyArray random key used to reset the environment. required Returns: Type Description state State object corresponding to the new state of the environment, timestep: TimeStep object corresponding the first timestep returned by the environment, step ( self , state : ~ State , action : Union [ jax . Array , numpy . ndarray , numpy . bool_ , numpy . number ]) -> Tuple [ ~ State , jumanji . types . TimeStep ] # Run one timestep of the environment's dynamics. Parameters: Name Type Description Default state ~State State object containing the dynamics of the environment. required action Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number] Array containing the action to take. required Returns: Type Description state State object corresponding to the next state of the environment, timestep: TimeStep object corresponding the timestep returned by the environment, observation_spec ( self ) -> Spec # Returns the observation spec. action_spec ( self ) -> Spec # Returns the action spec. render ( self , state : ~ State ) -> Any # Compute render frames during initialisation of the environment. Parameters: Name Type Description Default state ~State State object containing the dynamics of the environment. required close ( self ) -> None # Perform any necessary cleanup. Environments will automatically :meth: close() themselves when garbage collected or when the program exits. __enter__ ( self ) -> Wrapper special # __exit__ ( self , * args : Any ) -> None special # JumanjiToDMEnvWrapper ( Environment ) # A wrapper that converts Environment to dm_env.Environment. unwrapped : Environment property readonly # __init__ ( self , env : Environment , key : Optional [ jax . _src . prng . PRNGKeyArray ] = None ) special # Create the wrapped environment. Parameters: Name Type Description Default env Environment Environment to wrap to a dm_env.Environment . required key Optional[jax._src.prng.PRNGKeyArray] optional key to initialize the Environment with. None reset ( self ) -> TimeStep # Starts a new sequence and returns the first TimeStep of this sequence. Returns: Type Description A `TimeStep` namedtuple containing step_type: A StepType of FIRST . reward: None , indicating the reward is undefined. discount: None , indicating the discount is undefined. observation: A NumPy array, or a nested dict, list or tuple of arrays. Scalar values that can be cast to NumPy arrays (e.g. Python floats) are also valid in place of a scalar array. Must conform to the specification returned by observation_spec() . step ( self , action : ndarray ) -> TimeStep # Updates the environment according to the action and returns a TimeStep . If the environment returned a TimeStep with StepType.LAST at the previous step, this call to step will start a new sequence and action will be ignored. This method will also start a new sequence if called after the environment has been constructed and reset has not been called. Again, in this case action will be ignored. Parameters: Name Type Description Default action ndarray A NumPy array, or a nested dict, list or tuple of arrays corresponding to action_spec() . required Returns: Type Description A `TimeStep` namedtuple containing step_type: A StepType value. reward: Reward at this timestep, or None if step_type is StepType.FIRST . Must conform to the specification returned by reward_spec() . discount: A discount in the range [0, 1], or None if step_type is StepType.FIRST . Must conform to the specification returned by discount_spec() . observation: A NumPy array, or a nested dict, list or tuple of arrays. Scalar values that can be cast to NumPy arrays (e.g. Python floats) are also valid in place of a scalar array. Must conform to the specification returned by observation_spec() . observation_spec ( self ) -> Array # Returns the dm_env observation spec. action_spec ( self ) -> Array # Returns the dm_env action spec. MultiToSingleWrapper ( Wrapper ) # A wrapper that converts a multi-agent Environment to a single-agent Environment. __init__ ( self , env : Environment , reward_aggregator : Callable = < function sum at 0x7f157acde700 > , discount_aggregator : Callable = < function amax at 0x7f157acdeee0 > ) special # Create the wrapped environment. Parameters: Name Type Description Default env Environment Environment to wrap to a dm_env.Environment . required reward_aggregator Callable a function to aggregate all agents rewards into a single scalar value, e.g. sum.