pyRDDLGym-jax 2.2__py3-none-any.whl → 2.4__py3-none-any.whl

pyRDDLGym_jax/core/tuning.py

@@ -18,6 +18,7 @@ import datetime
 import threading
 import multiprocessing
 import os
+import termcolor
 import time
 import traceback
 from typing import Any, Callable, Dict, Iterable, Optional, Tuple
@@ -45,8 +46,7 @@ try:
     from pyRDDLGym_jax.core.visualization import JaxPlannerDashboard
 except Exception:
     raise_warning('Failed to load the dashboard visualization tool: '
-                  'please make sure you have installed the required packages.', 
-                  'red')
+                  'please make sure you have installed the required packages.', 'red')
     traceback.print_exc()
     JaxPlannerDashboard = None
 
@@ -159,24 +159,24 @@ class JaxParameterTuning:
         kernel3 = Matern(length_scale=5.0, length_scale_bounds=(1.0, 5.0), nu=2.5)
         return weight1 * kernel1 + weight2 * kernel2 + weight3 * kernel3
         
-    def summarize_hyperparameters(self) -> None:
+    def summarize_hyperparameters(self) -> str:
         hyper_params_table = []
         for (_, param) in self.hyperparams_dict.items():
             hyper_params_table.append(f'        {str(param)}')
         hyper_params_table = '\n'.join(hyper_params_table)
-        print(f'hyperparameter optimizer parameters:\n'
-              f'    tuned_hyper_parameters    =\n{hyper_params_table}\n'
-              f'    initialization_args       ={self.gp_init_kwargs}\n'
-              f'    gp_params                 ={self.gp_params}\n'
-              f'    tuning_iterations         ={self.gp_iters}\n'
-              f'    tuning_timeout            ={self.timeout_tuning}\n'
-              f'    tuning_batch_size         ={self.num_workers}\n'
-              f'    mp_pool_context_type      ={self.pool_context}\n'
-              f'    mp_pool_poll_frequency    ={self.poll_frequency}\n'
-              f'meta-objective parameters:\n'
-              f'    planning_trials_per_iter  ={self.eval_trials}\n'
-              f'    rollouts_per_trial        ={self.rollouts_per_trial}\n'
-              f'    acquisition_fn            ={self.acquisition}')
+        return (f'hyperparameter optimizer parameters:\n'
+                f'    tuned_hyper_parameters    =\n{hyper_params_table}\n'
+                f'    initialization_args       ={self.gp_init_kwargs}\n'
+                f'    gp_params                 ={self.gp_params}\n'
+                f'    tuning_iterations         ={self.gp_iters}\n'
+                f'    tuning_timeout            ={self.timeout_tuning}\n'
+                f'    tuning_batch_size         ={self.num_workers}\n'
+                f'    mp_pool_context_type      ={self.pool_context}\n'
+                f'    mp_pool_poll_frequency    ={self.poll_frequency}\n'
+                f'meta-objective parameters:\n'
+                f'    planning_trials_per_iter  ={self.eval_trials}\n'
+                f'    rollouts_per_trial        ={self.rollouts_per_trial}\n'
+                f'    acquisition_fn            ={self.acquisition}')
         
     @staticmethod
     def annealing_acquisition(n_samples: int, n_delay_samples: int=0,
@@ -346,6 +346,7 @@ class JaxParameterTuning:
         
         # remove keywords that should not be in the tuner
         train_args.pop('dashboard', None)
+        planner_args.pop('parallel_updates', None)
     
         # initialize env for evaluation (need fresh copy to avoid concurrency)
         env = RDDLEnv(domain, instance, vectorized=True, enforce_action_constraints=False)
@@ -368,12 +369,12 @@ class JaxParameterTuning:
     
     def tune_optimizer(self, optimizer: BayesianOptimization) -> None:
         '''Tunes the Bayesian optimization algorithm hyper-parameters.'''
-        print('\n' + f'The current kernel is {repr(optimizer._gp.kernel_)}.')
+        print(f'Kernel: {repr(optimizer._gp.kernel_)}.')
         
     def tune(self, key: int, log_file: str, show_dashboard: bool=False) -> ParameterValues:
         '''Tunes the hyper-parameters for Jax planner, returns the best found.'''
         
-        self.summarize_hyperparameters()
+        print(self.summarize_hyperparameters())
         
         # clear and prepare output file
         with open(log_file, 'w', newline='') as file:
@@ -445,13 +446,15 @@ class JaxParameterTuning:
                 # check if there is enough time left for another iteration
                 elapsed = time.time() - start_time
                 if elapsed >= self.timeout_tuning:
-                    print(f'global time limit reached at iteration {it}, aborting')
+                    message = termcolor.colored(
+                        f'[INFO] Global time limit reached at iteration {it}.', 'green')
+                    print(message)
                     break
                 
                 # continue with next iteration
                 print('\n' + '*' * 80 + 
                       f'\n[{datetime.timedelta(seconds=elapsed)}] ' + 
-                      f'starting iteration {it + 1}' + 
+                      f'Starting iteration {it + 1}' + 
                       '\n' + '*' * 80)
                 key, *subkeys = jax.random.split(key, num=num_workers + 1)
                 rows = [None] * num_workers
@@ -507,7 +510,10 @@ class JaxParameterTuning:
                 
                 # print best parameter if found
                 if best_target > old_best_target:
-                    print(f'* found new best average reward {best_target:.6f}')
+                    message = termcolor.colored(
+                        f'[INFO] Found new best average reward {best_target:.6f}.', 
+                        'green')
+                    print(message)
                 
                 # tune the optimizer here
                 self.tune_optimizer(optimizer)
@@ -528,7 +534,7 @@ class JaxParameterTuning:
         
         # print summary of results
         elapsed = time.time() - start_time
-        print(f'summary of hyper-parameter optimization:\n'
+        print(f'Summary of hyper-parameter optimization:\n'
               f'    time_elapsed         ={datetime.timedelta(seconds=elapsed)}\n'
               f'    iterations           ={it + 1}\n'
               f'    best_hyper_parameters={best_params}\n'

pyRDDLGym_jax/examples/run_plan.py

@@ -36,8 +36,8 @@ def main(domain, instance, method, episodes=1):
         abs_path = os.path.dirname(os.path.abspath(__file__))
         config_path = os.path.join(abs_path, 'configs', f'{domain}_{method}.cfg') 
         if not os.path.isfile(config_path):
-            raise_warning(f'Config file {config_path} was not found, '
-                          f'using default_{method}.cfg.', 'red')
+            raise_warning(f'[WARN] Config file {config_path} was not found, '
+                          f'using default_{method}.cfg.', 'yellow')
             config_path = os.path.join(abs_path, 'configs', f'default_{method}.cfg') 
     elif os.path.isfile(method):
         config_path = method

pyRDDLGym_jax/examples/run_scipy.py

@@ -31,8 +31,8 @@ def main(domain, instance, method, episodes=1):
     abs_path = os.path.dirname(os.path.abspath(__file__))
     config_path = os.path.join(abs_path, 'configs', f'{domain}_slp.cfg') 
     if not os.path.isfile(config_path):
-        raise_warning(f'Config file {config_path} was not found, '
-                      f'using default_slp.cfg.', 'red')
+        raise_warning(f'[WARN] Config file {config_path} was not found, '
+                      f'using default_slp.cfg.', 'yellow')
         config_path = os.path.join(abs_path, 'configs', 'default_slp.cfg') 
     planner_args, _, train_args = load_config(config_path)
 

{pyrddlgym_jax-2.2.dist-info → pyrddlgym_jax-2.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pyRDDLGym-jax
-Version: 2.2
+Version: 2.4
 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.2.dist-info → pyrddlgym_jax-2.4.dist-info}/RECORD

@@ -1,19 +1,19 @@
-pyRDDLGym_jax/__init__.py,sha256=lqo7WXKfZGHPIOxgE6EWI5fGZHP2h6XrwVNNVQAUN3Q,19
+pyRDDLGym_jax/__init__.py,sha256=6Bd43-94X_2dH_ErGLQ0_DvlhX5cLWkVPvn31JBzFkY,19
 pyRDDLGym_jax/entry_point.py,sha256=dxDlO_5gneEEViwkLCg30Z-KVzUgdRXaKuFjoZklkA0,974
 pyRDDLGym_jax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pyRDDLGym_jax/core/compiler.py,sha256=_ERueJW7GQ7S8-IezreeuLs3fNCZbQZ8j7VMUVlEt1k,82306
-pyRDDLGym_jax/core/logic.py,sha256=ZeCwCLqC6BvXpRT06TvE2bfPNO6ALuMzPmUvXNzW6Uw,52278
-pyRDDLGym_jax/core/planner.py,sha256=0rluBXKGNHRPEPfegOWcx9__cJHr8KjZdDJtG7i1JjI,122793
+pyRDDLGym_jax/core/compiler.py,sha256=NFWfTHtGf7F-t7Qhn6X-VpSAJkTVHm-oRjujFw4O1HA,82605
+pyRDDLGym_jax/core/logic.py,sha256=lfc2ak_ap_ajMEFlB5EHCRNgJym31dNyA-5d-7N4CZA,56271
+pyRDDLGym_jax/core/planner.py,sha256=wZJiZHV0Qxi9DS3AQ9Rx1doBvsKQXc1HYziY6GXTu_A,136965
 pyRDDLGym_jax/core/simulator.py,sha256=DnPL93WVCMZqtqMUoiJdfWcH9pEvNgGfDfO4NV0wIS0,9271
-pyRDDLGym_jax/core/tuning.py,sha256=RKKtDZp7unvfbhZEoaunZtcAn5xtzGYqXBB_Ij_Aapc,24205
+pyRDDLGym_jax/core/tuning.py,sha256=Gm3YJF84_2vDIIJpOj0tK0-4rlJoEjYwxRt_JpUKAOA,24482
 pyRDDLGym_jax/core/visualization.py,sha256=4BghMp8N7qtF0tdyDSqtxAxNfP9HPrQWTiXzAMJmx7o,70365
 pyRDDLGym_jax/core/assets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pyRDDLGym_jax/core/assets/favicon.ico,sha256=RMMrI9YvmF81TgYG7FO7UAre6WmYFkV3B2GmbA1l0kM,175085
 pyRDDLGym_jax/examples/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pyRDDLGym_jax/examples/run_gradient.py,sha256=KhXvijRDZ4V7N8NOI2WV8ePGpPna5_vnET61YwS7Tco,2919
 pyRDDLGym_jax/examples/run_gym.py,sha256=rXvNWkxe4jHllvbvU_EOMji_2-2k5d4tbBKhpMm_Gaw,1526
-pyRDDLGym_jax/examples/run_plan.py,sha256=v2AvwgIa4Ejr626vBOgWFJIQvay3IPKWno02ztIFCYc,2768
-pyRDDLGym_jax/examples/run_scipy.py,sha256=wvcpWCvdjvYHntO95a7JYfY2fuCMUTKnqjJikW0PnL4,2291
+pyRDDLGym_jax/examples/run_plan.py,sha256=TVfziHHaEC56wxwRw9llZ5iqSHe3m6yy8HxiR2TyvXE,2778
+pyRDDLGym_jax/examples/run_scipy.py,sha256=7uVnDXb7D3NTJqA2L8nrcYDJP-k0ba9dl9YqA2CD9ac,2301
 pyRDDLGym_jax/examples/run_tune.py,sha256=WbGO8RudIK-cPMAMKvI8NbFQAqkG-Blbnta3Efsep6c,3828
 pyRDDLGym_jax/examples/configs/Cartpole_Continuous_gym_drp.cfg,sha256=mE8MqhOlkHeXIGEVrnR3QY6I-_iy4uxFYRA71P1bmtk,347
 pyRDDLGym_jax/examples/configs/Cartpole_Continuous_gym_replan.cfg,sha256=nFFYHCKQUMn8x-OpJwu2pwe1tycNSJ8iAIwSkCBn33E,370
@@ -41,9 +41,9 @@ pyRDDLGym_jax/examples/configs/default_slp.cfg,sha256=mJo0woDevhQCSQfJg30ULVy9qG
 pyRDDLGym_jax/examples/configs/tuning_drp.cfg,sha256=CQMpSCKTkGioO7U82mHMsYWFRsutULx0V6Wrl3YzV2U,504
 pyRDDLGym_jax/examples/configs/tuning_replan.cfg,sha256=m_0nozFg_GVld0tGv92Xao_KONFJDq_vtiJKt5isqI8,501
 pyRDDLGym_jax/examples/configs/tuning_slp.cfg,sha256=KHu8II6CA-h_HblwvWHylNRjSvvGS3VHxN7JQNR4p_Q,464
-pyrddlgym_jax-2.2.dist-info/LICENSE,sha256=Y0Gi6H6mLOKN-oIKGZulQkoTJyPZeAaeuZu7FXH-meg,1095
-pyrddlgym_jax-2.2.dist-info/METADATA,sha256=aFNUX6uUZZHS7lPbYBTmMSH6TBiWmXbEgQNxPZNWiRI,17021
-pyrddlgym_jax-2.2.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-pyrddlgym_jax-2.2.dist-info/entry_points.txt,sha256=Q--z9QzqDBz1xjswPZ87PU-pib-WPXx44hUWAFoBGBA,59
-pyrddlgym_jax-2.2.dist-info/top_level.txt,sha256=n_oWkP_BoZK0VofvPKKmBZ3NPk86WFNvLhi1BktCbVQ,14
-pyrddlgym_jax-2.2.dist-info/RECORD,,
+pyrddlgym_jax-2.4.dist-info/LICENSE,sha256=Y0Gi6H6mLOKN-oIKGZulQkoTJyPZeAaeuZu7FXH-meg,1095
+pyrddlgym_jax-2.4.dist-info/METADATA,sha256=98Nl3EnEk-fRLeoy9orDScaikCT9M8X4zOfYtiS-WXI,17021
+pyrddlgym_jax-2.4.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
+pyrddlgym_jax-2.4.dist-info/entry_points.txt,sha256=Q--z9QzqDBz1xjswPZ87PU-pib-WPXx44hUWAFoBGBA,59
+pyrddlgym_jax-2.4.dist-info/top_level.txt,sha256=n_oWkP_BoZK0VofvPKKmBZ3NPk86WFNvLhi1BktCbVQ,14
+pyrddlgym_jax-2.4.dist-info/RECORD,,

{pyrddlgym_jax-2.2.dist-info → pyrddlgym_jax-2.4.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.8.2)
+Generator: setuptools (76.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any