rxnn 0.2.44__py3-none-any.whl → 0.2.45__py3-none-any.whl

rxnn/training/mrl.py

@@ -91,6 +91,7 @@ class MrlTrajectoryEpisode(TypedDict):
     reset_stm: bool
     steps: list[MrlTrajectoryStep]
 
+OptimField: TypeAlias = Literal['lr', 'critic_lr', 'weight_decay', 'critic_weight_decay', 'separate_memory_lr', 'memory_lr']
 
 class MRLTrainer:
     def __init__(
@@ -981,8 +982,19 @@ class MRLTrainer:
         self.reward = config.get('reward_model', self.shared_reward_model)  # MRL Reward Model for curriculum stage
         self.update_epochs = config.get('update_epochs', self.shared_update_epochs)  # Internal update epochs
         self.freeze_embeddings = config.get('freeze_embeddings', self.shared_freeze_embeddings)
-        if config['lr'] is not None or config['critic_lr'] is not None or config['weight_decay'] is not None or config[
-            'critic_weight_decay'] is not None or (config['separate_memory_lr'] and config['memory_lr'] is not None):
+
+
+
+        def has_param(field: OptimField) -> bool:
+            return field in config and config[field] is not None
+
+        optim_params: list[OptimField] = ['lr', 'critic_lr', 'weight_decay', 'critic_weight_decay']
+
+        has_any_optim_param = any(
+            has_param(field) for field in optim_params
+        ) or (has_param('separate_memory_lr') and config['separate_memory_lr'] and has_param('memory_lr'))
+
+        if has_any_optim_param:
             if config.get('separate_memory_lr', False):
                 self.optim_config = {
                     'lr': config.get('lr', self.base_optim_config['lr']),

{rxnn-0.2.44.dist-info → rxnn-0.2.45.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.2.44
+Version: 0.2.45
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning

{rxnn-0.2.44.dist-info → rxnn-0.2.45.dist-info}/RECORD

@@ -17,7 +17,7 @@ rxnn/training/callbacks.py,sha256=p72lbzFAmFjpcUvyy4aUB3qd53I8C6Sk5w9nQvsKgTk,35
 rxnn/training/dataset.py,sha256=7hTilFWPpqUEc6zNcMqBPjxFKxCfvTKKF3E8tVlwccQ,51250
 rxnn/training/ddp.py,sha256=VsNBjn3cY-uUj8hbsW7oKvb0_ZKnXnJ2KgObm-Mr9i4,836
 rxnn/training/models.py,sha256=tqABOt_xEcWbZNEW2I2Jt-3eyaGICK011zILwuTk6Zc,9082
-rxnn/training/mrl.py,sha256=BvrwqrIIyg_EmUA5p7c6UBcfFQ0ePIcl-EHEFQqyl2E,59472
+rxnn/training/mrl.py,sha256=L4G7xSPlxsymvNhvsSloCpaqYjOXxEm7GmKilM_Ojvc,59809
 rxnn/training/reward.py,sha256=B7nerPk9eNAv2i7umtNF88tVQVwijNNrchIrEITGHKk,11623
 rxnn/training/rl.py,sha256=q4NzIZAmXRHVToT13IHrPTtEikWQUvT0NO0IjApjAO8,6171
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
@@ -33,7 +33,7 @@ rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.44.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.44.dist-info/METADATA,sha256=tW2Ve4whRK2LfCxix10dTLS5Dl_0C6KhcK8FsoKq-x0,25960
-rxnn-0.2.44.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.44.dist-info/RECORD,,
+rxnn-0.2.45.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.45.dist-info/METADATA,sha256=g8IqXAR2tXEyaNQOs--IPFFtSOnrWe4oouPK1PQBITw,25960
+rxnn-0.2.45.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.45.dist-info/RECORD,,