Shared encoder DPNet

Author: Danfoa

data/linear_system/group=C10-dim=10/n_constraints=0/f_time_constant=1.1[s]-frames=200-horizon=6.4[s]/noise_level=0/n_trajs=118-train.pkl

@@ -55,7 +55,7 @@ def __init__(
             orth_w: float = 0.1,
             enforce_constant_fn: bool = True,
             use_spectral_score: bool = True,
-            aux_obs_space: bool = False,
+            shared_encoder: bool = True,
             obs_fn_params: Optional[dict] = None,
             linear_decoder: bool = True,
             **markov_dyn_params
@@ -67,7 +67,7 @@ def __init__(
         self.ck_w = ck_w
         self.orth_w = orth_w
         self.use_spectral_score = use_spectral_score
-        self.aux_obs_space = aux_obs_space
+        self.shared_encoder = shared_encoder
         self.inverse_projector = None  # if linear decoder is true, this is the map between obs to states.
         self.inverse_projector_bias = None
         self.linear_decoder = linear_decoder
@@ -169,7 +169,7 @@ def compute_loss_and_metrics(self,
         return loss, obs_space_metrics
 
     def get_obs_space_metrics(self, obs_state_traj: Tensor, obs_state_traj_aux: Optional[Tensor] = None) -> dict:
-        if obs_state_traj_aux is None and self.aux_obs_space:
+        if obs_state_traj_aux is None and self.shared_encoder:
             raise ValueError("aux_obs_space is True but obs_state_traj_aux is None")
         # Compute Covariance and Cross-Covariance operators for the observation state space.
         # Spectral and Projection scores, and CK loss terms.
@@ -270,14 +270,18 @@ def approximate_transfer_operator(self, train_data_loader: DataLoader):
     def build_obs_fn(self, num_layers, identity=False, **kwargs):
         if identity:
             return lambda x: (x, x)
-        obs_fn = MLP(in_dim=self.state_dim, out_dim=self.obs_state_dim, num_layers=num_layers,
-                     head_with_activation=False, **kwargs)
-        obs_fn_aux = None
-        if self.aux_obs_space:
-            obs_fn_aux = MLP(in_dim=self.state_dim, out_dim=self.obs_state_dim, num_layers=num_layers,
-                             head_with_activation=False, **kwargs)
-
-        return ObservableNet(obs_fn=obs_fn, obs_fn_aux=obs_fn_aux)
+        num_hidden_units = kwargs['num_hidden_units']
+        # Define the feature extractor used by the observable function.
+        encoder = MLP(in_dim=self.state_dim,
+                      out_dim=num_hidden_units,
+                      num_layers=num_layers,
+                      head_with_activation=True, **kwargs)
+        aux_encoder = None
+        if not self.shared_encoder:
+            aux_encoder = MLP(in_dim=self.state_dim, out_dim=num_hidden_units, num_layers=num_layers,
+                              head_with_activation=True, **kwargs)
+
+        return ObservableNet(encoder=encoder, aux_encoder=aux_encoder, obs_dim=self.obs_state_dim)
 
     def build_inv_obs_fn(self, num_layers, linear_decoder: bool, identity=False, **kwargs):
         if identity:

data/linear_system/group=C10-dim=10/n_constraints=0/f_time_constant=1.1[s]-frames=200-horizon=6.4[s]/noise_level=0/n_trajs=26-test.pkl

@@ -273,20 +273,20 @@ def build_obs_fn(self, num_layers, **kwargs):
                                   in_type=self.state_type_iso,
                                   desired_hidden_units=num_hidden_units)
 
-        obs_fn = EMLP(in_type=self.state_type_iso,
-                      out_type=self.obs_state_type,
-                      num_layers=num_layers,
-                      activation=act,
-                      **kwargs)
-        obs_fn_aux = None
-        if self.aux_obs_space:
-            obs_fn_aux = EMLP(in_type=self.state_type_iso,
-                              out_type=self.obs_state_type,
-                              num_layers=num_layers,
-                              activation=act,
-                              **kwargs)
-
-        return ObservableNet(obs_fn=obs_fn, obs_fn_aux=obs_fn_aux)
+        encoder = EMLP(in_type=self.state_type_iso,
+                       out_type=act.out_type,
+                       num_layers=num_layers,
+                       activation=act,
+                       **kwargs)
+        aux_encoder = None
+        if not self.shared_encoder:
+            aux_encoder = EMLP(in_type=self.state_type_iso,
+                               out_type=act.out_type,
+                               num_layers=num_layers,
+                               activation=act,
+                               **kwargs)
+
+        return ObservableNet(encoder=encoder, aux_encoder=aux_encoder, obs_type=self.obs_state_type)
 
     def build_inv_obs_fn(self, num_layers, linear_decoder: bool, **kwargs):
         if linear_decoder:

data/linear_system/group=C10-dim=10/n_constraints=0/f_time_constant=1.1[s]-frames=200-horizon=6.4[s]/noise_level=0/n_trajs=26-val.pkl

@@ -66,7 +66,7 @@ def training_step(self, batch, batch_idx):
 
         self.log("loss/train", loss, prog_bar=False)
         self.log_metrics(scalar_metrics, suffix="train", batch_size=self._batch_size)
-        # self.log_vector_metrics(vector_metrics, type_sufix="train", batch_size=self._batch_size)
+        self.log_vector_metrics(vector_metrics, type_sufix="train", batch_size=self._batch_size)
         return loss
 
     def validation_step(self, batch, batch_idx):
@@ -76,7 +76,7 @@ def validation_step(self, batch, batch_idx):
 
         self.log("loss/val", loss, prog_bar=False)
         self.log_metrics(scalar_metrics, suffix="val", batch_size=self._batch_size)
-        # self.log_vector_metrics(vector_metrics, type_sufix="val", batch_size=self._batch_size)
+        self.log_vector_metrics(vector_metrics, type_sufix="val", batch_size=self._batch_size)
         return {'output': outputs, 'input': batch}
 
     def test_step(self, batch, batch_idx):
@@ -117,7 +117,7 @@ def on_fit_start(self) -> None:
         self._loss_metrics_fn = loss_metrics_fn
 
     def on_train_start(self):
-        self.log("noise_level", self.trainer.datamodule.noise_level, prog_bar=False, on_epoch=True)
+        # self.log("noise_level", self.trainer.datamodule.noise_level, prog_bar=False, on_epoch=True)
 
         if hasattr(self.model, "approximate_transfer_operator"):
             metrics = self.model.approximate_transfer_operator(self.trainer.datamodule.predict_dataloader())
@@ -147,7 +147,7 @@ def on_train_end(self) -> None:
             self.compute_figure_metrics(self.val_metrics_fn, self.trainer.datamodule.train_dataloader(), suffix="train")
 
     def on_validation_start(self) -> None:
-        if hasattr(self.model, "approximate_transfer_operator"):
+        if hasattr(self.model, "approximate_transfer_operator") and self.trainer.current_epoch % 2 == 0:
             metrics = self.model.approximate_transfer_operator(self.trainer.datamodule.predict_dataloader())
             vector_metrics, scalar_metrics = self.separate_vector_scalar_metrics(metrics)
             self.log_metrics(scalar_metrics, suffix='')
@@ -192,13 +192,18 @@ def log_vector_metrics(self, metrics: Optional[dict]=None, type_sufix='', batch_
 
         for metric, vector in flat_metrics.items():
             assert vector.ndim >= 1, f"Vector metric {metric} has to be of shape (n_samples,) or (batch, time_steps)."
-            # Separate the last _sufix part from the key to obtain the metric name.
-            tmp = metric.split('_')   # Average value will use this name, vector metric will use the full name.
-            metric_name, metric_sufix = '_'.join(tmp[:-1]), tmp[-1]
-
             metric_log_name = f"{metric}/{type_sufix}"
-            # self.log(metric_log_name, torch.mean(vector), prog_bar=False, batch_size=batch_size)
+            if "_t/" in metric_log_name or "_dist/" in metric_log_name:
+                # Separate the last _sufix part from the key to obtain the metric name.
+                tmp = metric.split('_')   # Average value will use this name, vector metric will use the full name.
+                metric_name, metric_sufix = '_'.join(tmp[:-1]), tmp[-1]
+                self.log(f"{metric_name}/{type_sufix}", torch.mean(vector), prog_bar=False, batch_size=batch_size)
+            else:
+                self.log(metric_log_name, torch.mean(vector), prog_bar=False, batch_size=batch_size)
 
+            if type_sufix == 'train' or type_sufix == 'val':
+                continue
+                
             if metric_log_name in self._log_cache:
                 self._log_cache[metric_log_name] = np.concatenate([self._log_cache[metric_log_name], vector.detach().cpu().numpy()], axis=0)
             else:

data/linear_system/group=C10-dim=10/n_constraints=0/f_time_constant=1.1[s]-frames=200-horizon=6.4[s]/noise_level=1/n_trajs=118-train.pkl

@@ -5,44 +5,47 @@
 import torch.nn
 from escnn.nn import EquivariantModule
 
-from nn.EquivLinearDynamics import EquivLinearDynamics
-from nn.LinearDynamics import LinearDynamics
-from nn.mlp import MLP
-from nn.emlp import EMLP
-
-
 class ObservableNet(torch.nn.Module):
-
-    def __init__(self, 
-                 obs_fn: Union[torch.nn.Module, EquivariantModule],
-                 obs_fn_aux: Optional[Union[torch.nn.Module, EquivariantModule]] = None):
+    """
+    A network computing the observation state in the initial observable space H with measure μ(t) and the observable
+    space H' representing the observable space after a Δt step, which potentially has a different measure μ(t+Δt).
+    """
+    def __init__(self,
+                 encoder: Union[torch.nn.Module, EquivariantModule],
+                 aux_encoder: Optional[Union[torch.nn.Module, EquivariantModule]] = None,
+                 obs_dim: Optional[int] = None,
+                 obs_type: Optional[escnn.nn.FieldType] = None):
         super().__init__()
-        self.equivariant = isinstance(obs_fn, EquivariantModule)
-        self.use_aux_obs_fn = obs_fn_aux is not None
-
-        self.obs = obs_fn
-        self.obs_aux = None
-        if self.use_aux_obs_fn:  # Use two twin networks to compute the main and auxiliary observable space.
-            self.obs_aux = obs_fn_aux
+        self.equivariant = isinstance(encoder, EquivariantModule)
+        self.use_aux_encoder = aux_encoder is not None
+
+        self.encoder = encoder
+        self.aux_encoder = aux_encoder if self.use_aux_encoder else None
+
+        # Setting the bias of the linear layer to true is equivalent to setting the constant function in the basis
+        # of the space of functions. Then the bias of each dimension is the coefficient of the constant function.
+        if self.equivariant:
+            # Bias term (a.k.a the constant function) is present only on the trivial isotypic subspace
+            assert obs_type is not None, f"obs state Field type must be provided when using equivariant encoder"
+            self.obs_H = escnn.nn.Linear(
+                in_type=self.encoder.out_type, out_type=obs_type, bias=True)
+            self.obs_H_prime = escnn.nn.Linear(
+                in_type=self.encoder.out_type, out_type=obs_type, bias=True)
         else:
-            # Setting the bias of the linear layer to true is equivalent to setting the constant function in the basis
-            # of the space of functions. Then the bias of each dimension is the coefficient of the constant function.
-            if self.equivariant:
-                self.transfer_op_H_H_prime = escnn.nn.Linear(
-                    in_type=self.obs.out_type, out_type=self.obs.out_type, bias=True)
-            else:
-                self.transfer_op_H_H_prime = torch.nn.Linear(
-                    in_features=self.obs.out_dim, out_features=self.obs.out_dim, bias=True)
+            assert obs_dim is not None, f"obs state dimension must be provided when using non-equivariant encoder"
+            self.obs_H = torch.nn.Linear(
+                in_features=self.encoder.out_dim, out_features=obs_dim, bias=True)
+            self.obs_H_prime = torch.nn.Linear(
+                in_features=self.encoder.out_dim, out_features=obs_dim, bias=True)
 
     def forward(self, input):
 
-        obs_state = self.obs(input)
+        features = self.encoder(input)
+        aux_features = self.aux_encoder(input) if self.use_aux_encoder else features
 
-        if self.use_aux_obs_fn:
-            obs_aux_state = self.obs_aux(input)
-        else:
-            obs_aux_state = self.transfer_op_H_H_prime(obs_state)
+        obs_state_H = self.obs_H(features)
+        obs_state_H_prime = self.obs_H_prime(aux_features)
 
-        return obs_state, obs_aux_state
+        return obs_state_H, obs_state_H_prime
 
 

data/linear_system/group=C10-dim=10/n_constraints=0/f_time_constant=1.1[s]-frames=200-horizon=6.4[s]/noise_level=1/n_trajs=26-test.pkl

@@ -120,7 +120,7 @@ def get_activation(activation, in_type: FieldType, desired_hidden_units: int):
                                              irreps=list(unique_irreps),
                                              function=f"p_{activation.lower()}",
                                              inplace=True,
-                                             type='regular' if not group.continuous else 'rand',
+                                             type='rand', #'regular' if not group.continuous else 'rand'#TODO: fix this
                                              N=grid_length)
 
     def forward(self, x):