Equiv DPNets working-Iso basis bug fixed

Author: Danfoa

cfg/config.yaml

@@ -37,7 +37,7 @@ hydra:
     chdir: True
     env_set:
       XLA_PYTHON_CLIENT_PREALLOCATE: 'false'
-      HYDRA_FULL_ERROR: '1'
+      HYDRA_FULL_ERROR: 1
 #      CUDA_VISIBLE_DEVICES: ${.device}
     config:
       override_dirname:
@@ -54,6 +54,8 @@ hydra:
           - model.equivariant
           - model.max_ck_window_length
           - model.activation
+          - model.num_layers
+          - model.num_hidden_units
           - model.lr
           - model.batch_norm
           - model.bias

cfg/model/base_model.yaml

@@ -6,7 +6,7 @@ augment: False          # Weather to use data-augmentation of the dataset if the
 # Architectural and Task parameters
 equivariant: True      # Impose equivariance constraints on NNs
 activation: ReLU
-n_layers: 4             # Number MLPs' layers (including input and output layers)
+num_layers: 4             # Number MLPs' layers (including input and output layers)
 bias: True              # Use bias in the MLPs
 batch_norm: True        # Use batch normalization layers before activation functions in MLPs
 

cfg/model/dpnet.yaml

@@ -3,20 +3,25 @@ defaults:
 
 name: DPNet
 # Model hyperparameters
+
+# Symmetry exploitation parameters
 equivariant: False
-activation: Identity
-n_layers: 2             # Number MLPs' layers (including input and output layers)
-batch_norm: False       # Something wrong happens when we turn batch norm on. Performance goes to hell.
+group_avg_trick: True
+
+activation: ELU
+num_layers: 4             # Number MLPs' layers (including input and output layers)
+num_hidden_units: 128      # Number of hidden units in each layer
+batch_norm: ${model.equivariant}       # Something wrong happens when we turn batch norm on. Performance goes to hell.
 bias: False
 # Optimization hyperparameters parameters
 lr: 1e-3
 batch_size: 1024
 
 max_ck_window_length: ${system.pred_horizon}    # Maximum length of the Chapman-Kolmogorov window
-ck_w: 0.0                                   # Weight of the Chapman-Kolmogorov regularization term in the loss function
-orth_w: 0.1                                 # Weight of the orthonormal regularization term in the loss function
-aux_obs_space: False                     # Whether to use an auxiliary observable space.
-use_spectral_score: True                    # Whether to use the spectral or the correlation score
+ck_w: 0.0                                       # Weight of the Chapman-Kolmogorov regularization term in the loss function
+orth_w: 0.5                                     # Weight of the orthonormal regularization term in the loss function
+aux_obs_space: False                            # Whether to use an auxiliary observable space.
+use_spectral_score: True                        # Whether to use the spectral or the correlation score
 
-summary: ${model.name}-Equiv:${model.equivariant}-CK_w:${model.ck_w}-Orth_w:${model.orth_w}-Win:${model.max_ck_window_length}-Act:${model.activation}-B:${model.bias}-BN:${model.batch_norm}-LR:${model.lr}-L:${model.n_layers}-AOS:${model.aux_obs_space}-SS:${model.use_spectral_score}
+summary: ${model.name}-Equiv:${model.equivariant}-CK_w:${model.ck_w}-Orth_w:${model.orth_w}-Win:${model.max_ck_window_length}-Act:${model.activation}-B:${model.bias}-BN:${model.batch_norm}-LR:${model.lr}-L:${model.num_layers}-${model.num_hidden_units}-AOS:${model.aux_obs_space}-SS:${model.use_spectral_score}
 

cfg/system/linear_system.yaml

@@ -10,6 +10,10 @@ frames_per_state: 1          # Number of time-frames to use as a Markov Process
 pred_horizon: 10             # Number (or percentage) of Markov Process state time steps to predict into the future
 eval_pred_horizon: 100       # Number (or percentage) of Markov Process state time steps to predict into the future
 
-data_path: '${system.name}/${system.state_dim}-dim/'
+group: C3
+noise_level: 0
+n_constraints: 0
 
-summary: S:${system.state_dim}-OS:${system.obs_state_dim}-H:${system.pred_horizon}-Hv:${system.eval_pred_horizon}-F:${system.frames_per_state}
+data_path: '${system.name}/group=${system.group}-dim=${system.state_dim}/'
+
+summary: S:${system.state_dim}-OS:${system.obs_state_dim}-H:${system.pred_horizon}-G:${system.group}-N:${system.noise_level}

data/DynamicsDataModule.py

@@ -25,6 +25,7 @@ def __init__(self,
                  data_path: Path,
                  pred_horizon: Union[int, float] = 0.25,
                  eval_pred_horizon: Union[int, float] = 0.5,
+                 system_cfg: Optional[dict] = None,
                  batch_size: int = 256,
                  frames_per_step: int = 1,
                  num_workers: int = 0,
@@ -34,8 +35,11 @@ def __init__(self,
                  action_measurements: Optional[list[str]] = None,
                  ):
         super().__init__()
+        if system_cfg is None:
+            system_cfg = {}
         assert data_path.exists(), f"Data folder not found {data_path.absolute()}"
         self._data_path = data_path
+        self.system_cfg = system_cfg
         self.augment = augment
         self.frames_per_step = frames_per_step
         if isinstance(pred_horizon, float):
@@ -74,11 +78,12 @@ def prepare_data(self):
         start_time = time.time()
         log.info(f"Preparing datasets {self._data_path}")
 
-        path_to_dyn_sys_data = set([a.parent for a in list(self._data_path.rglob('*train.pkl'))])
-        # TODO: Handle multiple files from
-        system_data_path = path_to_dyn_sys_data.pop()
-        if len(path_to_dyn_sys_data) > 1:
-            raise NotImplementedError("Multiple dynamical systems not supported yet")
+        dyn_sys_data = set([a.parent for a in list(self._data_path.rglob('*train.pkl'))])
+        noise_level = self.system_cfg.get('noise_level', 0)
+        system_data_path = [path for path in dyn_sys_data if f"noise_level={noise_level}" in str(path)]
+        if len(system_data_path) > 1:
+            raise RuntimeError(f"Multiple potential paths {system_data_path} found")
+        system_data_path = system_data_path.pop()
 
         train_data, test_data, val_data = get_train_test_val_file_paths(system_data_path)
         # Obtain hugging face Iterable datasets instances