accumulate gradfix

experiments/custom-ppg.yaml

@@ -70,7 +70,7 @@ procgen-ppo:
         entropy_schedule: False
         
         # Memory management, if batch size overflow, batch splitting is done to handle it
-        max_minibatch_size: 1500
+        max_minibatch_size: 1000
         updates_per_batch: 8
 
         normalize_actions: False

experiments/custom-torch-ppo.yaml

@@ -61,7 +61,7 @@ procgen-ppo:
         entropy_schedule: False
         
         # Memory management, if batch size overflow, batch splitting is done to handle it
-        max_minibatch_size: 1500
+        max_minibatch_size: 1000
         updates_per_batch: 8
 
         normalize_actions: False
@@ -77,7 +77,7 @@ procgen-ppo:
         # === Settings for Model ===
         model:
             custom_model: impala_torch_custom
-            custom_options: 
+            custom_model_config: 
                 depths: [32, 64, 64]
                 nlatents: 512
                 use_layernorm: True

models/impala_torch_custom.py

@@ -87,11 +87,11 @@ class ImpalaCNN(TorchModelV2, nn.Module):
         nn.Module.__init__(self)
 
         self.device = device
-        depths = model_config['custom_options'].get('depths') or [16, 32, 32]
-        nlatents = model_config['custom_options'].get('nlatents') or 256
-        d2rl = model_config['custom_options'].get('d2rl') or False
-        self.use_layernorm = model_config['custom_options'].get('use_layernorm') or False
-        self.diff_framestack = model_config['custom_options'].get('diff_framestack') or False
+        depths = model_config['custom_model_config'].get('depths')
+        nlatents = model_config['custom_model_config'].get('nlatents')
+        d2rl = model_config['custom_model_config'].get('d2rl')
+        self.use_layernorm = model_config['custom_model_config'].get('use_layernorm')
+        self.diff_framestack = model_config['custom_model_config'].get('diff_framestack')
 
         
         h, w, c = obs_space.shape