pyRDDLGym-jax 1.2__tar.gz → 2.0__tar.gz

{pyrddlgym_jax-1.2 → pyrddlgym_jax-2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pyRDDLGym-jax
-Version: 1.2
+Version: 2.0
 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

pyrddlgym_jax-2.0/pyRDDLGym_jax/__init__.py

@@ -0,0 +1 @@
+__version__ = '2.0'

{pyrddlgym_jax-1.2 → pyrddlgym_jax-2.0}/pyRDDLGym_jax/core/compiler.py

@@ -1,3 +1,18 @@
+# ***********************************************************************
+# JAXPLAN
+# 
+# Author: Michael Gimelfarb
+#
+# REFERENCES:
+# 
+# [1] Gimelfarb, Michael, Ayal Taitler, and Scott Sanner. "JaxPlan and GurobiPlan: 
+# Optimization Baselines for Replanning in Discrete and Mixed Discrete-Continuous 
+# Probabilistic Domains." Proceedings of the International Conference on Automated 
+# Planning and Scheduling. Vol. 34. 2024.
+#
+# ***********************************************************************
+
+
 from functools import partial
 import traceback
 from typing import Any, Callable, Dict, List, Optional
@@ -524,7 +539,7 @@ class JaxRDDLCompiler:
                 _jax_wrapped_single_step_policy,
                 in_axes=(0, None, None, None, 0, None)
             )(keys, policy_params, hyperparams, step, subs, model_params)
-            model_params = jax.tree_map(lambda x: jnp.mean(x, axis=0), model_params)
+            model_params = jax.tree_map(partial(jnp.mean, axis=0), model_params)
             carry = (key, policy_params, hyperparams, subs, model_params)
             return carry, log            
             

{pyrddlgym_jax-1.2 → pyrddlgym_jax-2.0}/pyRDDLGym_jax/core/logic.py

@@ -1,4 +1,31 @@
-from typing import Optional, Set
+# ***********************************************************************
+# JAXPLAN
+# 
+# Author: Michael Gimelfarb
+#
+# REFERENCES:
+#
+# [1] Gimelfarb, Michael, Ayal Taitler, and Scott Sanner. "JaxPlan and GurobiPlan: 
+# Optimization Baselines for Replanning in Discrete and Mixed Discrete-Continuous 
+# Probabilistic Domains." Proceedings of the International Conference on Automated 
+# Planning and Scheduling. Vol. 34. 2024.
+#
+# [2] Petersen, Felix, Christian Borgelt, Hilde Kuehne, and Oliver Deussen. "Learning with 
+# algorithmic supervision via continuous relaxations." Advances in Neural Information 
+# Processing Systems 34 (2021): 16520-16531.
+#
+# [3] Agustsson, Eirikur, and Lucas Theis. "Universally quantized neural compression." 
+# Advances in neural information processing systems 33 (2020): 12367-12376.
+#
+# [4] Gupta, Madan M., and J11043360726 Qi. "Theory of T-norms and fuzzy inference 
+# methods." Fuzzy sets and systems 40, no. 3 (1991): 431-450.
+#
+# [5] Jang, Eric, Shixiang Gu, and Ben Poole. "Categorical Reparametrization with 
+# Gumble-Softmax." In International Conference on Learning Representations (ICLR 2017). 
+# OpenReview. net, 2017.
+#
+# ***********************************************************************
+
 
 import jax
 import jax.numpy as jnp
@@ -759,14 +786,14 @@ class FuzzyLogic(Logic):
     
     def __str__(self) -> str:
         return (f'model relaxation:\n'
-                f'    tnorm         ={str(self.tnorm)}\n'
-                f'    complement    ={str(self.complement)}\n'
-                f'    comparison    ={str(self.comparison)}\n'
-                f'    sampling      ={str(self.sampling)}\n'
-                f'    rounding      ={str(self.rounding)}\n'
-                f'    control       ={str(self.control)}\n'
-                f'    underflow_tol ={self.eps}\n'
-                f'    use_64_bit    ={self.use64bit}')
+                f'    tnorm        ={str(self.tnorm)}\n'
+                f'    complement   ={str(self.complement)}\n'
+                f'    comparison   ={str(self.comparison)}\n'
+                f'    sampling     ={str(self.sampling)}\n'
+                f'    rounding     ={str(self.rounding)}\n'
+                f'    control      ={str(self.control)}\n'
+                f'    underflow_tol={self.eps}\n'
+                f'    use_64_bit   ={self.use64bit}\n')
 
     def summarize_hyperparameters(self) -> None:
         print(self.__str__())