pyRDDLGym-jax 2.6__tar.gz → 2.8__tar.gz

{pyrddlgym_jax-2.6 → pyrddlgym_jax-2.8}/LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2024 pyrddlgym-project
+Copyright (c) 2025 pyrddlgym-project
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

{pyrddlgym_jax-2.6 → pyrddlgym_jax-2.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyRDDLGym-jax
-Version: 2.6
+Version: 2.8
 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
@@ -20,7 +20,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: pyRDDLGym>=2.3
+Requires-Dist: pyRDDLGym>=2.5
 Requires-Dist: tqdm>=4.66
 Requires-Dist: jax>=0.4.12
 Requires-Dist: optax>=0.1.9

pyrddlgym_jax-2.8/pyRDDLGym_jax/__init__.py

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

{pyrddlgym_jax-2.6 → pyrddlgym_jax-2.8}/pyRDDLGym_jax/core/compiler.py

@@ -30,7 +30,8 @@ from pyRDDLGym.core.debug.exception import (
     print_stack_trace, 
     raise_warning,
     RDDLInvalidNumberOfArgumentsError, 
-    RDDLNotImplementedError
+    RDDLNotImplementedError,
+    RDDLUndefinedVariableError
 )
 from pyRDDLGym.core.debug.logger import Logger
 from pyRDDLGym.core.simulator import RDDLSimulatorPrecompiled
@@ -56,7 +57,8 @@ class JaxRDDLCompiler:
                  allow_synchronous_state: bool=True,
                  logger: Optional[Logger]=None,
                  use64bit: bool=False,
-                 compile_non_fluent_exact: bool=True) -> None:
+                 compile_non_fluent_exact: bool=True,
+                 python_functions: Optional[Dict[str, Callable]]=None) -> None:
         '''Creates a new RDDL to Jax compiler.
         
         :param rddl: the RDDL model to compile into Jax
@@ -65,7 +67,8 @@ class JaxRDDLCompiler:
         :param logger: to log information about compilation to file
         :param use64bit: whether to use 64 bit arithmetic
         :param compile_non_fluent_exact: whether non-fluent expressions 
-        are always compiled using exact JAX expressions.
+        are always compiled using exact JAX expressions
+        :param python_functions: dictionary of external Python functions to call from RDDL
         '''
         self.rddl = rddl
         self.logger = logger
@@ -99,11 +102,15 @@ class JaxRDDLCompiler:
         self.traced = tracer.trace()
         
         # extract the box constraints on actions
+        if python_functions is None:
+            python_functions = {}
+        self.python_functions = python_functions
         simulator = RDDLSimulatorPrecompiled(
             rddl=self.rddl,
             init_values=self.init_values,
             levels=self.levels,
-            trace_info=self.traced
+            trace_info=self.traced,
+            python_functions=python_functions
         )  
         constraints = RDDLConstraints(simulator, vectorized=True)
         self.constraints = constraints
@@ -605,6 +612,8 @@ class JaxRDDLCompiler:
             jax_expr = self._jax_aggregation(expr, init_params)
         elif etype == 'func':
             jax_expr = self._jax_functional(expr, init_params)
+        elif etype == 'pyfunc':
+            jax_expr = self._jax_pyfunc(expr, init_params)
         elif etype == 'control':
             jax_expr = self._jax_control(expr, init_params)
         elif etype == 'randomvar':
@@ -926,6 +935,84 @@ class JaxRDDLCompiler:
         raise RDDLNotImplementedError(
             f'Function {op} is not supported.\n' + print_stack_trace(expr))   
     
+    def _jax_pyfunc(self, expr, init_params):
+        NORMAL = JaxRDDLCompiler.ERROR_CODES['NORMAL']
+
+        # get the Python function by name
+        _, pyfunc_name = expr.etype
+        pyfunc = self.python_functions.get(pyfunc_name)
+        if pyfunc is None:
+            raise RDDLUndefinedVariableError(
+                f'Undefined external Python function <{pyfunc_name}>, '
+                f'must be one of {list(self.python_functions.keys())}.\n' +  
+                print_stack_trace(expr))
+        
+        captured_vars, args = expr.args
+        scope_vars = self.traced.cached_objects_in_scope(expr)
+        dest_indices = self.traced.cached_sim_info(expr)
+        free_vars = [p for p in scope_vars if p[0] not in captured_vars]
+        free_dims = self.rddl.object_counts(p for (_, p) in free_vars)
+        num_free_vars = len(free_vars)
+        captured_types = [t for (p, t) in scope_vars if p in captured_vars]
+        require_dims = self.rddl.object_counts(captured_types)
+
+        # compile the inputs to the function
+        jax_inputs = [self._jax(arg, init_params) for arg in args]
+
+        # compile the function evaluation function
+        def _jax_wrapped_external_function(x, params, key):
+
+            # evaluate inputs to the function
+            # first dimensions are non-captured vars in outer scope followed by all the _
+            error = NORMAL
+            flat_samples = []
+            for jax_expr in jax_inputs:
+                sample, key, err, params = jax_expr(x, params, key)
+                shape = jnp.shape(sample)
+                first_dim = 1
+                for dim in shape[:num_free_vars]:
+                    first_dim *= dim
+                new_shape = (first_dim,) + shape[num_free_vars:]
+                flat_sample = jnp.reshape(sample, new_shape)
+                flat_samples.append(flat_sample)
+                error |= err
+
+            # now all the inputs have dimensions equal to (k,) + the number of _ occurences
+            # k is the number of possible non-captured object combinations
+            # evaluate the function independently for each combination
+            # output dimension for each combination is captured variables (n1, n2, ...)
+            # so the total dimension of the output array is (k, n1, n2, ...)
+            sample = jax.vmap(pyfunc, in_axes=0)(*flat_samples)
+            if not isinstance(sample, jnp.ndarray):
+                raise ValueError(
+                    f'Output of external Python function <{pyfunc_name}> '
+                    f'is not a JAX array.\n' + print_stack_trace(expr))                        
+            
+            pyfunc_dims = jnp.shape(sample)[1:]
+            if len(require_dims) != len(pyfunc_dims):
+                raise ValueError(
+                    f'External Python function <{pyfunc_name}> returned array with '
+                    f'{len(pyfunc_dims)} dimensions, which does not match the '
+                    f'number of captured parameter(s) {len(require_dims)}.\n' +  
+                    print_stack_trace(expr))
+            for (param, require_dim, actual_dim) in zip(captured_vars, require_dims, pyfunc_dims):
+                if require_dim != actual_dim:
+                    raise ValueError(
+                        f'External Python function <{pyfunc_name}> returned array with '
+                        f'{actual_dim} elements for captured parameter <{param}>, '
+                        f'which does not match the number of objects {require_dim}.\n' + 
+                        print_stack_trace(expr))
+
+            # unravel the combinations k back into their original dimensions
+            sample = jnp.reshape(sample, free_dims + pyfunc_dims)
+            
+            # rearrange the output dimensions to match the outer scope
+            source_indices = [num_free_vars + i for i in range(len(pyfunc_dims))]
+            sample = jnp.moveaxis(sample, source=source_indices, destination=dest_indices)
+            return sample, key, error, params
+        
+        return _jax_wrapped_external_function
+
     # ===========================================================================
     # control flow
     # ===========================================================================

{pyrddlgym_jax-2.6 → pyrddlgym_jax-2.8}/pyRDDLGym_jax/core/planner.py

@@ -1810,7 +1810,8 @@ class JaxBackpropPlanner:
                  dashboard_viz: Optional[Any]=None,
                  print_warnings: bool=True,
                  parallel_updates: Optional[int]=None,
-                 preprocessor: Optional[Preprocessor]=None) -> None:
+                 preprocessor: Optional[Preprocessor]=None,
+                 python_functions: Optional[Dict[str, Callable]]=None) -> None:
         '''Creates a new gradient-based algorithm for optimizing action sequences
         (plan) in the given RDDL. Some operations will be converted to their
         differentiable counterparts; the specific operations can be customized
@@ -1853,6 +1854,7 @@ class JaxBackpropPlanner:
         :param print_warnings: whether to print warnings
         :param parallel_updates: how many optimizers to run independently in parallel
         :param preprocessor: optional preprocessor for state inputs to plan
+        :param python_functions: dictionary of external Python functions to call from RDDL
         '''
         self.rddl = rddl
         self.plan = plan
@@ -1879,7 +1881,10 @@ class JaxBackpropPlanner:
         self.use_pgpe = pgpe is not None
         self.print_warnings = print_warnings
         self.preprocessor = preprocessor
-        
+        if python_functions is None:
+            python_functions = {}
+        self.python_functions = python_functions
+
         # set optimizer
         try:
             optimizer = optax.inject_hyperparams(optimizer)(**optimizer_kwargs)
@@ -2027,7 +2032,8 @@ r"""
             use64bit=self.use64bit,
             cpfs_without_grad=self.cpfs_without_grad,
             compile_non_fluent_exact=self.compile_non_fluent_exact,
-            print_warnings=self.print_warnings
+            print_warnings=self.print_warnings,
+            python_functions=self.python_functions
         )
         self.compiled.compile(log_jax_expr=True, heading='RELAXED MODEL')
         
@@ -2035,7 +2041,8 @@ r"""
         self.test_compiled = JaxRDDLCompiler(
             rddl=rddl,
             logger=self.logger,
-            use64bit=self.use64bit
+            use64bit=self.use64bit,
+            python_functions=self.python_functions
         )
         self.test_compiled.compile(log_jax_expr=True, heading='EXACT MODEL')
         

{pyrddlgym_jax-2.6 → pyrddlgym_jax-2.8}/pyRDDLGym_jax/core/simulator.py

@@ -20,7 +20,7 @@
 
 import time
 import numpy as np
-from typing import Dict, Optional, Union
+from typing import Callable, Dict, Optional, Union
 
 import jax
 
@@ -48,6 +48,7 @@ class JaxRDDLSimulator(RDDLSimulator):
                  logger: Optional[Logger]=None,
                  keep_tensors: bool=False,
                  objects_as_strings: bool=True,
+                 python_functions: Optional[Dict[str, Callable]]=None,
                  **compiler_args) -> None:
         '''Creates a new simulator for the given RDDL model with Jax as a backend.
         
@@ -60,8 +61,9 @@ class JaxRDDLSimulator(RDDLSimulator):
         :param logger: to log information about compilation to file
         :param keep_tensors: whether the sampler takes actions and
         returns state in numpy array form
-        param objects_as_strings: whether to return object values as strings (defaults
+        :param objects_as_strings: whether to return object values as strings (defaults
         to integer indices if False)
+        :param python_functions: dictionary of external Python functions to call from RDDL
         :param **compiler_args: keyword arguments to pass to the Jax compiler
         '''
         if key is None:
@@ -73,7 +75,8 @@ class JaxRDDLSimulator(RDDLSimulator):
         # generate direct sampling with default numpy RNG and operations
         super(JaxRDDLSimulator, self).__init__(
             rddl, logger=logger, 
-            keep_tensors=keep_tensors, objects_as_strings=objects_as_strings)
+            keep_tensors=keep_tensors, objects_as_strings=objects_as_strings,
+            python_functions=python_functions)
     
     def seed(self, seed: int) -> None:
         super(JaxRDDLSimulator, self).seed(seed)
@@ -83,7 +86,12 @@ class JaxRDDLSimulator(RDDLSimulator):
         rddl = self.rddl
         
         # compilation
-        compiled = JaxRDDLCompiler(rddl, logger=self.logger, **self.compiler_args)
+        compiled = JaxRDDLCompiler(
+            rddl, 
+            logger=self.logger, 
+            python_functions=self.python_functions, 
+            **self.compiler_args
+        )
         compiled.compile(log_jax_expr=True, heading='SIMULATION MODEL')
         
         self.init_values = compiled.init_values

{pyrddlgym_jax-2.6 → pyrddlgym_jax-2.8}/pyRDDLGym_jax.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyRDDLGym-jax
-Version: 2.6
+Version: 2.8
 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
@@ -20,7 +20,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: pyRDDLGym>=2.3
+Requires-Dist: pyRDDLGym>=2.5
 Requires-Dist: tqdm>=4.66
 Requires-Dist: jax>=0.4.12
 Requires-Dist: optax>=0.1.9

{pyrddlgym_jax-2.6 → pyrddlgym_jax-2.8}/pyRDDLGym_jax.egg-info/requires.txt

@@ -1,4 +1,4 @@
-pyRDDLGym>=2.3
+pyRDDLGym>=2.5
 tqdm>=4.66
 jax>=0.4.12
 optax>=0.1.9

{pyrddlgym_jax-2.6 → pyrddlgym_jax-2.8}/setup.py

@@ -19,7 +19,7 @@ long_description = (Path(__file__).parent / "README.md").read_text()
 
 setup(
       name='pyRDDLGym-jax',
-      version='2.6',
+      version='2.8',
       author="Michael Gimelfarb, Ayal Taitler, Scott Sanner",
       author_email="mike.gimelfarb@mail.utoronto.ca, ataitler@gmail.com, ssanner@mie.utoronto.ca",
       description="pyRDDLGym-jax: automatic differentiation for solving sequential planning problems in JAX.",
@@ -29,7 +29,7 @@ setup(
       url="https://github.com/pyrddlgym-project/pyRDDLGym-jax",
       packages=find_packages(),
       install_requires=[
-          'pyRDDLGym>=2.3',
+          'pyRDDLGym>=2.5',
           'tqdm>=4.66',
           'jax>=0.4.12',
           'optax>=0.1.9',

pyrddlgym_jax-2.6/pyRDDLGym_jax/__init__.py

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