pyRDDLGym-jax 0.5__tar.gz → 1.1__tar.gz

{pyrddlgym_jax-0.5 → pyrddlgym_jax-1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pyRDDLGym-jax
-Version: 0.5
+Version: 1.1
 Summary: pyRDDLGym-jax: automatic differentiation for solving sequential planning problems in JAX.
 Home-page: https://github.com/pyrddlgym-project/pyRDDLGym-jax
 Author: Michael Gimelfarb, Ayal Taitler, Scott Sanner
@@ -25,51 +25,77 @@ Requires-Dist: tensorflow-probability>=0.21.0
 Provides-Extra: extra
 Requires-Dist: bayesian-optimization>=2.0.0; extra == "extra"
 Requires-Dist: rddlrepository>=2.0; extra == "extra"
+Provides-Extra: dashboard
+Requires-Dist: dash>=2.18.0; extra == "dashboard"
+Requires-Dist: dash-bootstrap-components>=1.6.0; extra == "dashboard"
 
 # pyRDDLGym-jax
 
-Author: [Mike Gimelfarb](https://mike-gimelfarb.github.io)
+![Python Version](https://img.shields.io/badge/python-3.9%2B-blue)
+[![PyPI Version](https://img.shields.io/pypi/v/pyRDDLGym-jax.svg)](https://pypi.org/project/pyRDDLGym-jax/)
+[![Documentation Status](https://readthedocs.org/projects/pyrddlgym/badge/?version=latest)](https://pyrddlgym.readthedocs.io/en/latest/jax.html)
+![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)
+[![Cumulative PyPI Downloads](https://img.shields.io/pypi/dm/pyrddlgym-jax)](https://pypistats.org/packages/pyrddlgym-jax)
 
-This directory provides:
-1. automated translation and compilation of RDDL description files into [JAX](https://github.com/google/jax), converting any RDDL domain to a differentiable simulator!
-2. powerful, fast and scalable gradient-based planning algorithms, with extendible and flexible policy class representations, automatic model relaxations for working in discrete and hybrid domains, and much more!
+[Installation](#installation) | [Run cmd](#running-from-the-command-line) | [Run python](#running-from-another-python-application) | [Configuration](#configuring-the-planner) | [Dashboard](#jaxplan-dashboard) | [Tuning](#tuning-the-planner) | [Simulation](#simulation) | [Citing](#citing-jaxplan)
 
-> [!NOTE]  
-> While Jax planners can support some discrete state/action problems through model relaxations, on some discrete problems it can perform poorly (though there is an ongoing effort to remedy this!).
-> If you find it is not making sufficient progress, check out the [PROST planner](https://github.com/pyrddlgym-project/pyRDDLGym-prost) (for discrete spaces) or the [deep reinforcement learning wrappers](https://github.com/pyrddlgym-project/pyRDDLGym-rl).
+**pyRDDLGym-jax (known in the literature as JaxPlan) is an efficient gradient-based/differentiable planning algorithm in JAX.** 
+
+Purpose:
+
+1. automatic translation of any RDDL description file into a differentiable simulator in JAX
+2. flexible policy class representations, automatic model relaxations for working in discrete and hybrid domains, and Bayesian hyper-parameter tuning.
+
+Some demos of solved problems by JaxPlan:
+
+<p align="middle">
+<img src="Images/intruders.gif" width="120" height="120" margin=0/>
+<img src="Images/marsrover.gif" width="120" height="120" margin=0/>
+<img src="Images/pong.gif" width="120" height="120" margin=0/>
+<img src="Images/quadcopter.gif" width="120" height="120" margin=0/>
+<img src="Images/reacher.gif" width="120" height="120" margin=0/>
+<img src="Images/reservoir.gif" width="120" height="120" margin=0/>
+</p>
 
-## Contents
+> [!WARNING]  
+> Starting in version 1.0 (major release), the ``weight`` parameter in the config file was removed, 
+and was moved to the individual logic components which have their own unique weight parameter assigned.
+> Furthermore, the tuning module has been redesigned from the ground up, and supports tuning arbitrary hyper-parameters via config templates!
+> Finally, the terrible visualizer for the planner was removed and replaced with an interactive real-time dashboard (similar to tensorboard, but custom designed for the planner)!
 
-- [Installation](#installation)
-- [Running from the Command Line](#running-from-the-command-line)
-- [Running from within Python](#running-from-within-python)
-- [Configuring the Planner](#configuring-the-planner)
-- [Simulation](#simulation)
-- [Manual Gradient Calculation](#manual-gradient-calculation)
-- [Citing pyRDDLGym-jax](#citing-pyrddlgym-jax)
+> [!NOTE]  
+> While JaxPlan can support some discrete state/action problems through model relaxations, on some discrete problems it can perform poorly (though there is an ongoing effort to remedy this!).
+> If you find it is not making sufficient progress, check out the [PROST planner](https://github.com/pyrddlgym-project/pyRDDLGym-prost) (for discrete spaces) or the [deep reinforcement learning wrappers](https://github.com/pyrddlgym-project/pyRDDLGym-rl).
   
 ## Installation
 
-To use the compiler or planner without the automated hyper-parameter tuning, you will need the following packages installed: 
-- ``pyRDDLGym>=2.0``
-- ``tqdm>=4.66``
-- ``jax>=0.4.12``
-- ``optax>=0.1.9``
-- ``dm-haiku>=0.0.10``
-- ``tensorflow-probability>=0.21.0``
+To install the bare-bones version of JaxPlan with **minimum installation requirements**:
 
-Additionally, if you wish to run the examples, you need ``rddlrepository>=2``. 
-To run the automated tuning optimization, you will also need ``bayesian-optimization>=2.0.0``.
+```shell
+pip install pyRDDLGym-jax
+```
 
-You can install pyRDDLGym-jax with all requirements using pip:
+To install JaxPlan with the **automatic hyper-parameter tuning** and rddlrepository:
 
 ```shell
 pip install pyRDDLGym-jax[extra]
 ```
 
+(Since version 1.0) To install JaxPlan with the **visualization dashboard**:
+
+```shell
+pip install pyRDDLGym-jax[dashboard]
+```
+
+(Since version 1.0) To install JaxPlan with **all options**:
+
+```shell
+pip install pyRDDLGym-jax[extra,dashboard]
+```
+
 ## Running from the Command Line
 
-A basic run script is provided to run the Jax Planner on any domain in ``rddlrepository`` from the install directory of pyRDDLGym-jax:
+A basic run script is provided to run JaxPlan on any domain in ``rddlrepository`` from the install directory of pyRDDLGym-jax:
 
 ```shell
 python -m pyRDDLGym_jax.examples.run_plan <domain> <instance> <method> <episodes>
@@ -86,35 +112,15 @@ The ``method`` parameter supports three possible modes:
 - ``drp`` is the deep reactive policy network described [in this paper](https://ojs.aaai.org/index.php/AAAI/article/view/4744)
 - ``replan`` is the same as ``slp`` except the plan is recalculated at every decision time step.
    
-A basic run script is also provided to run the automatic hyper-parameter tuning:
-
-```shell
-python -m pyRDDLGym_jax.examples.run_tune <domain> <instance> <method> <trials> <iters> <workers>
-```
-
-where:
-- ``domain`` is the domain identifier as specified in rddlrepository
-- ``instance`` is the instance identifier
-- ``method`` is the planning method to use (i.e. drp, slp, replan)
-- ``trials`` is the (optional) number of trials/episodes to average in evaluating each hyper-parameter setting
-- ``iters`` is the (optional) maximum number of iterations/evaluations of Bayesian optimization to perform
-- ``workers`` is the (optional) number of parallel evaluations to be done at each iteration, e.g. the total evaluations = ``iters * workers``.
-
-For example, the following will train the Jax Planner on the Quadcopter domain with 4 drones:
+For example, the following will train JaxPlan on the Quadcopter domain with 4 drones:
 
 ```shell
 python -m pyRDDLGym_jax.examples.run_plan Quadcopter 1 slp
 ```
 
-After several minutes of optimization, you should get a visualization as follows:
+## Running from Another Python Application
 
-<p align="center">
-<img src="Images/quadcopter.gif" width="400" height="400" margin=1/>
-</p>
-
-## Running from within Python
-
-To run the Jax planner from within a Python application, refer to the following example:
+To run JaxPlan from within a Python application, refer to the following example:
 
 ```python
 import pyRDDLGym
@@ -144,17 +150,15 @@ The basic structure of a configuration file is provided below for a straight-lin
 ```ini
 [Model]
 logic='FuzzyLogic'
-logic_kwargs={'weight': 20}
-tnorm='ProductTNorm'
-tnorm_kwargs={}
+comparison_kwargs={'weight': 20}
+rounding_kwargs={'weight': 20}
+control_kwargs={'weight': 20}
 
 [Optimizer]
 method='JaxStraightLinePlan'
 method_kwargs={}
 optimizer='rmsprop'
 optimizer_kwargs={'learning_rate': 0.001}
-batch_size_train=1
-batch_size_test=1
 
 [Training]
 key=42
@@ -179,7 +183,7 @@ method_kwargs={'topology': [128, 64], 'activation': 'tanh'}
 ```
 
 The configuration file must then be passed to the planner during initialization. 
-For example, the [previous script here](#running-from-within-python) can be modified to set parameters from a config file:
+For example, the [previous script here](#running-from-another-python-application) can be modified to set parameters from a config file:
 
 ```python
 from pyRDDLGym_jax.core.planner import load_config
@@ -193,6 +197,101 @@ controller = JaxOfflineController(planner, **train_args)
 ...
 ```
 
+### JaxPlan Dashboard
+
+Since version 1.0, JaxPlan has an optional dashboard that allows keeping track of the planner performance across multiple runs, 
+and visualization of the policy or model, and other useful debugging features.
+
+<p align="middle">
+<img src="Images/dashboard.png" width="480" height="248" margin=0/>
+</p>
+
+To run the dashboard, add the following entry to your config file:
+
+```ini
+...
+[Training]
+dashboard=True
+...
+```
+
+More documentation about this and other new features will be coming soon.
+
+### Tuning the Planner
+
+It is easy to tune the planner's hyper-parameters efficiently and automatically using Bayesian optimization.
+To do this, first create a config file template with patterns replacing concrete parameter values that you want to tune, e.g.:
+
+```ini
+[Model]
+logic='FuzzyLogic'
+comparison_kwargs={'weight': TUNABLE_WEIGHT}
+rounding_kwargs={'weight': TUNABLE_WEIGHT}
+control_kwargs={'weight': TUNABLE_WEIGHT}
+
+[Optimizer]
+method='JaxStraightLinePlan'
+method_kwargs={}
+optimizer='rmsprop'
+optimizer_kwargs={'learning_rate': TUNABLE_LEARNING_RATE}
+
+[Training]
+train_seconds=30
+print_summary=False
+print_progress=False
+train_on_reset=True
+```
+
+would allow to tune the sharpness of model relaxations, and the learning rate of the optimizer.
+
+Next, you must link the patterns in the config with concrete hyper-parameter ranges the tuner will understand:
+
+```python
+import pyRDDLGym
+from pyRDDLGym_jax.core.tuning import JaxParameterTuning, Hyperparameter
+
+# set up the environment   
+env = pyRDDLGym.make(domain, instance, vectorized=True)
+    
+# load the config file template with planner settings
+with open('path/to/config.cfg', 'r') as file: 
+    config_template = file.read() 
+    
+# map parameters in the config that will be tuned
+def power_10(x):
+    return 10.0 ** x
+    
+hyperparams = [
+    Hyperparameter('TUNABLE_WEIGHT', -1., 5., power_10),  # tune weight from 10^-1 ... 10^5
+    Hyperparameter('TUNABLE_LEARNING_RATE', -5., 1., power_10),   # tune lr from 10^-5 ... 10^1
+]
+    
+# build the tuner and tune
+tuning = JaxParameterTuning(env=env,
+                            config_template=config_template,
+                            hyperparams=hyperparams,
+                            online=False,
+                            eval_trials=trials,
+                            num_workers=workers,
+                            gp_iters=iters)
+tuning.tune(key=42, log_file='path/to/log.csv')
+```
+ 
+A basic run script is provided to run the automatic hyper-parameter tuning for the most sensitive parameters of JaxPlan:
+
+```shell
+python -m pyRDDLGym_jax.examples.run_tune <domain> <instance> <method> <trials> <iters> <workers>
+```
+
+where:
+- ``domain`` is the domain identifier as specified in rddlrepository
+- ``instance`` is the instance identifier
+- ``method`` is the planning method to use (i.e. drp, slp, replan)
+- ``trials`` is the (optional) number of trials/episodes to average in evaluating each hyper-parameter setting
+- ``iters`` is the (optional) maximum number of iterations/evaluations of Bayesian optimization to perform
+- ``workers`` is the (optional) number of parallel evaluations to be done at each iteration, e.g. the total evaluations = ``iters * workers``.
+ 
+
 ## Simulation
 
 The JAX compiler can be used as a backend for simulating and evaluating RDDL environments:
@@ -206,47 +305,17 @@ from pyRDDLGym_jax.core.simulator import JaxRDDLSimulator
 env = pyRDDLGym.make("domain", "instance", backend=JaxRDDLSimulator)
 
 # evaluate the random policy
-agent = RandomAgent(action_space=env.action_space,
-                    num_actions=env.max_allowed_actions)
+agent = RandomAgent(action_space=env.action_space, num_actions=env.max_allowed_actions)
 agent.evaluate(env, verbose=True, render=True)
 ```
 
 For some domains, the JAX backend could perform better than the numpy-based one, due to various compiler optimizations. 
 In any event, the simulation results using the JAX backend should (almost) always match the numpy backend.
 
-## Manual Gradient Calculation
-
-For custom applications, it is desirable to compute gradients of the model that can be optimized downstream. 
-Fortunately, we provide a very convenient function for compiling the transition/step function ``P(s, a, s')`` of the environment into JAX.
 
-```python
-import pyRDDLGym
-from pyRDDLGym_jax.core.planner import JaxRDDLCompilerWithGrad
+## Citing JaxPlan
 
-# set up the environment
-env = pyRDDLGym.make("domain", "instance", vectorized=True)
-
-# create the step function
-compiled = JaxRDDLCompilerWithGrad(rddl=env.model)
-compiled.compile()
-step_fn = compiled.compile_transition()
-```
-
-This will return a JAX compiled (pure) function requiring the following inputs:
-- ``key`` is the ``jax.random.PRNGKey`` key for reproducible randomness
-- ``actions`` is the dictionary of action fluent tensors
-- ``subs`` is the dictionary of state-fluent and non-fluent tensors
-- ``model_params`` are the parameters of the differentiable relaxations, such as ``weight``
-
-The function returns a dictionary containing a variety of variables, such as updated pvariables including next-state fluents (``pvar``), reward obtained (``reward``), error codes (``error``).
-It is thus possible to apply any JAX transformation to the output of the function, such as computing gradient using ``jax.grad()`` or batched simulation using ``jax.vmap()``.
-
-Compilation of entire rollouts is also possible by calling the ``compile_rollouts`` function.
-An [example is provided to illustrate how you can define your own policy class and compute the return gradient manually](https://github.com/pyrddlgym-project/pyRDDLGym-jax/blob/main/pyRDDLGym_jax/examples/run_gradient.py).
-
-## Citing pyRDDLGym-jax
-
-The [following citation](https://ojs.aaai.org/index.php/ICAPS/article/view/31480) describes the main ideas of the framework. Please cite it if you found it useful:
+The [following citation](https://ojs.aaai.org/index.php/ICAPS/article/view/31480) describes the main ideas of JaxPlan. Please cite it if you found it useful:
 
 ```
 @inproceedings{gimelfarb2024jaxplan,
@@ -258,21 +327,8 @@ The [following citation](https://ojs.aaai.org/index.php/ICAPS/article/view/31480
 }
 ```
 
-The utility optimization is discussed in [this paper](https://ojs.aaai.org/index.php/AAAI/article/view/21226):
-
-```
-@inproceedings{patton2022distributional,
-    title={A distributional framework for risk-sensitive end-to-end planning in continuous mdps},
-    author={Patton, Noah and Jeong, Jihwan and Gimelfarb, Mike and Sanner, Scott},
-    booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
-    volume={36},
-    number={9},
-    pages={9894--9901},
-    year={2022}
-}
-```
-
 Some of the implementation details derive from the following literature, which you may wish to also cite in your research papers:
+- [A Distributional Framework for Risk-Sensitive End-to-End Planning in Continuous MDPs](https://ojs.aaai.org/index.php/AAAI/article/view/21226)
 - [Deep reactive policies for planning in stochastic nonlinear domains, AAAI 2019](https://ojs.aaai.org/index.php/AAAI/article/view/4744)
 - [Scalable planning with tensorflow for hybrid nonlinear domains, NeurIPS 2017](https://proceedings.neurips.cc/paper/2017/file/98b17f068d5d9b7668e19fb8ae470841-Paper.pdf)