Mini Cheetah Sample Efficiency Experiments

Author: Danfoa

cfg/config.yaml

@@ -3,8 +3,6 @@ defaults:
   - model: dpnet
   - system: linear_system
   - override hydra/launcher: joblib
-#  - override hydra/sweeper: optuna
-#  - override hydra/sweeper/sampler: tpe
 
 # Running hyperparameters
 exp_name: test
@@ -76,27 +74,26 @@ hydra:
     dir: ./experiments/${hydra.job.name}/
     subdir: ${system.summary}_${model.summary}_${hydra.job.override_dirname}
 
-  job_logging:
-    version: 1
-    colorlog: True
-    formatters:
-      simple:
-        format: '[%(levelname)s][%(name)s] %(message)s'
+#  job_logging:
+#    version: 1
+#    colorlog: True
+#    root:
+#      handlers: [console, file_error]
+#      level: INFO
 #    handlers:
 #      console:
 #        class: logging.StreamHandler
 #        formatter: simple
 #        stream: ext://sys.stdout
-#      file:
-#        class: logging.handlers.RotatingFileHandler
+#      file_error:
+#        class: logging.FileHandler
+#        level: ERROR
 #        formatter: simple
-#        filename: log.log
-#        maxBytes: 2048
-#        backupCount: 1
-#    root:
-#      handlers: [console]
-#
-#    disable_existing_loggers: false
+#        filename: err.log
+#        mode: a
+#    formatters:
+#      simple:
+#        format: '[%(levelname)s][%(name)s] %(message)s'
 
 #  sweeper:
 ##    sampler:

cfg/model/dae-aug.yaml

@@ -0,0 +1,6 @@
+defaults:
+  - dae
+
+name: DAE-AUG
+augment: True
+

cfg/system/linear_system.yaml

@@ -17,7 +17,7 @@ pred_horizon: 2                                # Number (or percentage) of Marko
 eval_pred_horizon: 20                          # Number (or percentage) of Markov Process state time steps to predict into the future
 test_pred_horizon: .5                          # Number (or percentage) of Markov Process state time steps to predict into the future
 noise_level: 2                                 # Scale of the Wiener process noise from 0 to 9 (0 is no noise)
-n_constraints: 0
+n_constraints: 1
 
 # Symmetry parameters
 group: C3

cfg/system/mini_cheetah.yaml

@@ -5,22 +5,22 @@ name: 'mini_cheetah'
 subgroup_id: null
 
 max_epochs: 100
-early_stop_epochs: 20
+early_stop_epochs: 30
 standardize: True
 
-obs_state_ratio: 2                             # obs_state_dim = obs_state_ratio * state_dim
+obs_state_ratio: 5                             # obs_state_dim = obs_state_ratio * state_dim
 
-state_obs: ['q_js', 'v_js']
+state_obs: ['joint_pos', 'joint_vel', 'base_z', 'base_vel', 'base_ori', 'base_ang_vel']
 action_obs: []
 
 # dt = 0.0012 s   this is the average delta time between observations.
-frames_per_state: 1                            # Number of time-frames to use as a Markov Process state time step
-pred_horizon: 10                               # Number (or percentage) of Markov Process state time steps to predict into the future
-eval_pred_horizon: 20                          # Number (or percentage) of Markov Process state time steps to predict into the future
-test_pred_horizon: 200
+frames_per_state:  1                                 # Number of time-frames to use as a Markov Process state time step
+pred_horizon:      30                                # Number (or percentage) of Markov Process state time steps to predict into the future
+eval_pred_horizon: ${system.pred_horizon}            # Number (or percentage) of Markov Process state time steps to predict into the future
+test_pred_horizon: 150
 
-dynamic_mode: concrete_galloping
+dynamic_mode: forward_minus_0_4
 
-data_path: '${system.name}/recordings/${system.dynamic_mode}'
+data_path: '${system.name}/raysim_recordings/flat/${system.dynamic_mode}'
 
 summary: S:${system.dynamic_mode}-OS:${system.obs_state_ratio}-H:${system.pred_horizon}-EH:${system.eval_pred_horizon}

data/DynamicsDataModule.py

@@ -13,7 +13,7 @@
 from torch.utils.data import DataLoader, default_collate
 from tqdm import tqdm
 
-from data.DynamicsRecording import DynamicsRecording, get_dynamics_dataset, get_train_test_val_file_paths
+from morpho_symm.data.DynamicsRecording import DynamicsRecording, get_dynamics_dataset, get_train_test_val_file_paths
 from utils.mysc import traj_from_states
 from utils.plotting import plot_system_3D, plot_trajectories, plot_two_panel_trajectories
 
@@ -220,7 +220,7 @@ def compute_loss_metrics(self, predictions: dict, inputs: dict) -> (torch.Tensor
     def train_dataloader(self):
         return DataLoader(dataset=self.train_dataset, batch_size=self.batch_size,
                           num_workers=self.num_workers,
-                          collate_fn=self.collate_fn,
+                          collate_fn=self.collate_fn if not self.augment else self.data_augmentation_collate_fn,
                           persistent_workers=True if self.num_workers > 0 else False, drop_last=False)
 
     def val_dataloader(self):
@@ -257,6 +257,8 @@ def collate_fn(self, batch_list: list) -> dict:
 
     def data_augmentation_collate_fn(self, batch_list: list) -> dict:
         batch = torch.utils.data.default_collate(batch_list)
+        for k, v in batch.items():
+            batch[k] = v.to(self.device)
         state = batch['state']
         next_state = batch['next_state']
 
@@ -307,7 +309,7 @@ def plot_sample_trajs(self):
     assert path_to_data.exists(), f"Invalid Dataset path {path_to_data.absolute()}"
 
     # Find all dynamic systems recordings
-    path_to_data /= Path('mini_cheetah') / 'recordings' / 'grass'
+    path_to_data /= Path('mini_cheetah') / 'raysim_recordings' / 'flat' / 'forward_minus_0_4'
     # path_to_data = Path('/home/danfoa/Projects/koopman_robotics/data/linear_system/group=C10-dim=10/n_constraints=0/'
     #                     'f_time_constant=1.5[s]-frames=200-horizon=8.7[s]/noise_level=0')
     path_to_dyn_sys_data = set([a.parent for a in list(path_to_data.rglob('*train.pkl'))])
@@ -316,14 +318,14 @@ def plot_sample_trajs(self):
 
     data_module = DynamicsDataModule(data_path=mock_path,
                                      pred_horizon=10,
-                                     eval_pred_horizon=200,
-                                     test_pred_horizon=300,
+                                     eval_pred_horizon=10,
+                                     test_pred_horizon=150,
                                      frames_per_step=1,
                                      num_workers=1,
                                      batch_size=1000,
                                      augment=False,
                                      standardize=True,
-                                     state_obs=('q_js', 'v_js', 'imu_ang_vel'),
+                                     # state_obs=('base_z_error', 'base_vel_error', 'base_ang_vel_error', ),
                                      # action_obs=tuple(),
                                      )
 

data/DynamicsRecording.py

@@ -1,3 +1,4 @@
+import copy
 import logging
 import math
 import pickle
@@ -35,9 +36,13 @@ class DynamicsRecording:
     def save_to_file(self, file_path: Path):
         # Store representations and groups without serializing
         if len(self.obs_representations) > 0:
-            self._obs_rep_irreps = {k: rep.irreps for k, rep in self.obs_representations.items()}
-            self._obs_rep_names = {k: rep.name for k, rep in self.obs_representations.items()}
-            self._obs_rep_Q = {k: rep.change_of_basis for k, rep in self.obs_representations.items()}
+            self._obs_rep_irreps = {}
+            self._obs_rep_names = {}
+            self._obs_rep_Q = {}
+            for k, rep in self.obs_representations.items():
+                self._obs_rep_irreps[k] = rep.irreps if rep is not None else None
+                self._obs_rep_names[k] = rep.name if rep is not None else None
+                self._obs_rep_Q[k] = rep.change_of_basis if rep is not None else None
             group = self.obs_representations[self.state_obs[0]].group
             self._group_keys = group._keys
             self._group_name = group.__class__.__name__
@@ -335,7 +340,7 @@ def reduce_dataset_size(recordings: Iterable[DynamicsRecording], train_ratio: fl
         return recordings
     log.info(f"Reducing dataset size to {train_ratio * 100}%")
     # Ensure all training seeds use the same training data partitions
-    from utils.mysc import TemporaryNumpySeed
+    from morpho_symm.utils.mysc import TemporaryNumpySeed
     with TemporaryNumpySeed(10):
         for r in recordings:
             # Decide to keep a ratio of the original trajectories
@@ -380,6 +385,54 @@ def reduce_dataset_size(recordings: Iterable[DynamicsRecording], train_ratio: fl
                 new_recordings = {k: v[idx_to_keep] for k, v in r.recordings.items()}
                 r.recordings = new_recordings
 
+def split_train_val_test(
+        dyn_recording: DynamicsRecording, partition_sizes=(0.70, 0.15, 0.15)) -> tuple[DynamicsRecording]:
+    assert np.isclose(np.sum(partition_sizes), 1.0), f"Invalid partition sizes {partition_sizes}"
+    partitions_names = ["train", "val", "test"]
+
+    log.info(f"Partitioning {dyn_recording.description} into train/val/test of sizes {partition_sizes}[%]")
+    # Ensure all training seeds use the same training data partitions
+    from morpho_symm.utils.mysc import TemporaryNumpySeed
+    with TemporaryNumpySeed(10):   # Ensure deterministic behavior
+        # Decide to keep a ratio of the original trajectories
+        num_trajs = int(dyn_recording.info['num_traj'])
+        if num_trajs < 10:  # Do not discard entire trajectories, but rather parts of the trajectories
+            # Take the time horizon from the first observation
+            sample_obs = dyn_recording.recordings[dyn_recording.state_obs[0]]
+            if len(sample_obs.shape) == 3:  # [traj, time, obs_dim]
+                time_horizon = sample_obs.shape[1]
+            elif len(sample_obs.shape) == 2:  # [traj, obs_dim]
+                time_horizon = sample_obs.shape[0]
+            else:
+                raise RuntimeError(f"Invalid shape {sample_obs.shape} of {dyn_recording.state_obs[0]}")
+
+            num_samples = time_horizon
+            min_idx = 0
+            partitions_sample_idx = {partition: None for partition in partitions_names}
+            for partition_name, ratio in zip(partitions_names, partition_sizes):
+                max_idx = min_idx + int(num_samples * ratio)
+                partitions_sample_idx[partition_name] = list(range(min_idx, max_idx))
+                min_idx = min_idx + int(num_samples * ratio)
+
+            # TODO: Avoid deep copying the data itself.
+            partitions_recordings = {partition: copy.deepcopy(dyn_recording) for partition in partitions_names}
+            for partition_name, sample_idx in partitions_sample_idx.items():
+                part_num_samples = len(sample_idx)
+                partitions_recordings[partition_name].info['trajectory_length'] = part_num_samples
+                partitions_recordings[partition_name].recordings = dict()
+                for obs_name in dyn_recording.recordings.keys():
+                    if len(dyn_recording.recordings[obs_name].shape) == 3:
+                        data = dyn_recording.recordings[obs_name][:, sample_idx]
+                    elif len(dyn_recording.recordings[obs_name].shape) == 2:
+                        data = dyn_recording.recordings[obs_name][sample_idx]
+                    else:
+                        raise RuntimeError(f"Invalid shape {dyn_recording.recordings[obs_name].shape} of {obs_name}")
+                    partitions_recordings[partition_name].recordings[obs_name] = data
+
+            return partitions_recordings['train'], partitions_recordings['val'], partitions_recordings['test']
+        else:  # Discard entire trajectories
+            raise NotImplementedError()
+
 
 def get_dynamics_dataset(train_shards: list[Path],
                          test_shards: Optional[list[Path]] = None,

data/linear_system/stable_linear_system.py

@@ -326,7 +326,7 @@ def evolve_linear_dynamics(A: np.ndarray, init_state: np.ndarray, dt: float, sim
 if __name__ == '__main__':
     np.set_printoptions(precision=3)
     seed_everything(120)
-    order = 3
+    order = 5
     subgroups_ids = dict(C2=('cone', 1),
                          Tetrahedral=('fulltetra',),
                          Octahedral=(True, 'octa',),
@@ -347,7 +347,7 @@ def evolve_linear_dynamics(A: np.ndarray, init_state: np.ndarray, dt: float, sim
     # Parameters of the state space.
     for n_constraints in [1]: # [0, 1]:
         # Define the state representation.
-        for multiplicity in [2]:
+        for multiplicity in [20]:
             # Generate stable equivariant linear dynamics withing a range of fast and slow dynamics
             max_time_constant = 5  # [s] Maximum time constant of the system.
             min_period = max_time_constant / 3  # [s] Minimum period of oscillation of the fastest transient mode.

data/mini_cheetah/harmonic_analysis.py

@@ -99,7 +99,7 @@ def decom_signal_into_isotypic_components2(signal: np.ndarray, rep: Representati
     train_data = train_data[0]  # Get the first file path
     dyn_recording = DynamicsRecording.load_from_file(train_data)
 
-    robot, G = load_symmetric_system(robot_name='mini_cheetah')
+    robot, G = load_symmetric_system(robot_name='mini_cheetah-c2')
     q0 = robot._q0
 
     rep_Qjs = G.representations['Q_js']
@@ -125,18 +125,32 @@ def decom_signal_into_isotypic_components2(signal: np.ndarray, rep: Representati
     tau_js_t_iso_signal, tau_js_t_iso_comps = decom_signal_into_isotypic_components(tau_js_t, rep_TqQjs)
 
     # Compute approximate work done by the torques on the system
-    work_t = v_js_t * tau_js_t
+    work_t = np.sum(v_js_t * tau_js_t, axis=-1)
     work_iso_decomp_t = {}  # Decomposition of the work into work done by isotypic components forces.
     for irrep_id, v_js_t_iso in v_js_t_iso_signal.items():
         tau_js_t_iso = tau_js_t_iso_signal[irrep_id]
         work_iso_dims = v_js_t_iso * tau_js_t_iso
-        work_iso_t = v_js_t_iso * tau_js_t_iso
+        work_iso_t = np.sum(v_js_t_iso * tau_js_t_iso, axis=-1)
         work_iso_decomp_t[irrep_id] = work_iso_t
         # Check that computing the work in the original basis is equivalent to computing the work in the isotypic basis
         # v_js_t_iso_comp = v_js_t_iso_comps[irrep_id]
         # tau_js_t_iso_comp = tau_js_t_iso_comps[irrep_id]
         # assert np.allclose(work_iso_t, np.sum(v_js_t_iso_comp * tau_js_t_iso_comp, axis=-1))
 
+    work_iso_trajs = np.concatenate([work_iso[None, :, None] for work_iso in work_iso_decomp_t.values()], axis=0)
+    
+    fig = plot_trajectories(trajs=work_iso_trajs,
+                            # secondary_trajs=work[None, :, None],
+                            main_legend_label="Work_Iso",
+                            colorscale='G10',
+                            plot_error=False)
+
+    # fig = plot_trajectories(trajs=work_t[None, :, None],
+    #                         main_legend_label="Work",
+    #                         colorscale='Set2',
+    #                         plot_error=False,
+    #                         fig=fig)
+
     # Compute the Kinetic Energy of the system and the kinetic energy of each isotypic component.
     from pinocchio import pinocchio_pywrap as pin
     kin_energy_iso_t = {}
@@ -161,20 +175,30 @@ def decom_signal_into_isotypic_components2(signal: np.ndarray, rep: Representati
     import plotly.express as px
     df = px.data.gapminder()
 
-    work_iso_trajs = np.concatenate([work_iso[None, :, None] for work_iso in work_iso_decomp_t.values()], axis=0)
     kin_energy_iso_trajs = np.concatenate([kin_energy_iso[None, :, None] for kin_energy_iso in kin_energy_iso_t.values()], axis=0)
-    fig = plot_trajectories(trajs=kin_energy_iso_trajs,
-                            # secondary_trajs=work[None, :, None],
-                            main_legend_label="KinEnergy_Iso",
-                            colorscale='G10',
-                            plot_error=False)
+    # fig = plot_trajectories(trajs=kin_energy_iso_trajs,
+    #                         # secondary_trajs=work[None, :, None],
+    #                         main_legend_label="Work_Iso",
+    #                         colorscale='G10',
+    #                         plot_error=False)
 
-    # fig = plot_trajectories(trajs=work_t[None, :600, None],
+    # fig = plot_trajectories(trajs=work_t,
     #                         main_legend_label="Work",
     #                         colorscale='Set2',
     #                         plot_error=False,
     #                         fig=fig)
 
+    # fig = plot_trajectories(trajs=work_iso_trajs,
+    #                         # secondary_trajs=work[None, :, None],
+    #                         main_legend_label="Work_Iso",
+    #                         colorscale='G10',
+    #                         plot_error=False)
+    #
+    # fig = plot_trajectories(trajs=work_t,
+    #                         main_legend_label="Work",
+    #                         colorscale='Set2',
+    #                         plot_error=False,
+    #                         fig=fig)
 
     fig.show()
     # Plot the data

data/mini_cheetah/raysim_recordings/flat/forward_minus_0_4/heightmap_logger/state_reduced_nmpc.npy

@@ -91,7 +91,7 @@ def mat_to_dynamics_recordings(data_path: Path,
                 dynamics_parameters=dict(dt=dt, group=dict(group_name=G.name, group_order=G.order())),
                 recordings=dict(q_js=q_js_unit_circle_t[None, idx].astype(np.float32),
                                 v_js=v_js_t[None, idx].astype(np.float32),
-                                torques=torques[None, idx].astype(np.float32),
+                                joint_torques=torques[None, idx].astype(np.float32),
                                 # imu_acc=imu_acc[None, idx].astype(np.float32),
                                 imu_ori=imu_ori[None, idx].astype(np.float32),
                                 imu_ang_vel=imu_ang_vel[None, idx].astype(np.float32),
@@ -102,7 +102,7 @@ def mat_to_dynamics_recordings(data_path: Path,
                 action_obs=('torques',),
                 obs_representations=dict(q_js=rep_Q_js,
                                          v_js=rep_TqQ_js,
-                                         torques=rep_TqQ_js,
+                                         joint_torques=rep_TqQ_js,
                                          # imu_acc=rep_imu_acc,
                                          imu_ori=rep_imu_ori,
                                          imu_ang_vel=rep_imu_ang_vel,

data/mini_cheetah/raysim_recordings/flat/forward_minus_0_4_yawrate_0_4/heightmap_logger/state_reduced_nmpc.npy

@@ -1,6 +1,6 @@
 #!/bin/bash
-#PBS -l select=1:ncpus=16:mpiprocs=16:ngpus=2
-#PBS -l walltime=12:00:00
+#PBS -l select=1:ncpus=20:mpiprocs=20:ngpus=2
+#PBS -l walltime=24:00:00
 #PBS -N dae_obs_state_dim
 #PBS -M [email protected]
 #PBS -m bea
@@ -10,6 +10,8 @@
 cd /work/dordonez/Projects/koopman_robotics
 conda activate robotics
 
-python train_observables.py --multirun exp_name=C10-ObsStateDim hydra.launcher.n_jobs=-1 model=dae system.n_constraints=1 system.group=C10 system.state_dim=30 system.obs_state_ratio=1,4 system.train_ratio=1.0 seed=0,1,2,3 device=0 &
-python train_observables.py --multirun exp_name=C10-ObsStateDim hydra.launcher.n_jobs=-1 model=dae system.n_constraints=1 system.group=C10 system.state_dim=30 system.obs_state_ratio=2,3 system.train_ratio=1.0 seed=0,1,2,3 device=0 &
+#python train_observables.py --multirun seed=0,1,2,3 device=0 exp_name=C10-ObsStateDim2 hydra.launcher.n_jobs=-1 model=dae system.n_constraints=1 system.group=C10 system.state_dim=30 system.obs_state_ratio=1,4 system.train_ratio=1.0 &
+#python train_observables.py --multirun seed=0,1,2,3 device=1 exp_name=C10-ObsStateDim2 hydra.launcher.n_jobs=4 model=dae system.n_constraints=1 system.group=C10 system.state_dim=30 system.obs_state_ratio=2,3 system.train_ratio=1.0  &
+python train_observables.py --multirun seed=0,1,2,3 device=0 exp_name=C10-ObsStateDim2 hydra.launcher.n_jobs=-1 model=dae system.n_constraints=1 system.group=C5 system.state_dim=30 system.obs_state_ratio=1,4 system.train_ratio=1.0 &
+python train_observables.py --multirun seed=0,1,2,3 device=1 exp_name=C10-ObsStateDim2 hydra.launcher.n_jobs=4 model=dae system.n_constraints=1 system.group=C5 system.state_dim=30 system.obs_state_ratio=2,3 system.train_ratio=1.0  &
 wait