Group data also over game phase

Author: alexschmander

tools/machine-learning/fall_detection/data_loading/load_data.py

@@ -1,12 +1,14 @@
+from datetime import datetime
+from itertools import batched
 from pathlib import Path
 from typing import Any
-from mcap.reader import make_reader, McapReader
+
 import polars as pl
+from mcap.reader import McapReader, make_reader
 from msgpack import unpackb
-from datetime import datetime
-from itertools import batched
-from tqdm import tqdm
 from scipy.spatial.transform import Rotation
+from tqdm import tqdm
+
 from .unnest_structs import unnest_column
 
 
@@ -111,13 +113,15 @@ def iter_mcap(reader: McapReader, topics: list[str]):
 
 def read_mcap(mcap_path: Path) -> pl.DataFrame:
     robot_identifier = mcap_path.parts[-3]
+    game_phase_identifier = mcap_path.parts[-4]
     match_identifier = mcap_path.parts[-5]
 
     print(mcap_path)
     with open(mcap_path, "rb") as mcap_data:
         reader = make_reader(mcap_data)
         dataframe = pl.from_dicts(iter_mcap(reader, OUTPUTS)).with_columns(
             pl.lit(robot_identifier).alias("robot_identifier"),
+            pl.lit(game_phase_identifier).alias("game_phase_identifier"),
             pl.lit(match_identifier).alias("match_identifier"),
         )
     return dataframe
@@ -139,7 +143,7 @@ def convert_mcaps(mcaps: list[str]) -> pl.DataFrame:
 def load(path: str):
     df = pl.read_parquet(path).with_columns(
         (pl.col("time") - pl.col("time").min())
-        .over("robot_identifier", "match_identifier")
+        .over("robot_identifier", "game_phase_identifier", "match_identifier")
         .dt.total_seconds()
         .alias("time_in_game"),
     )

tools/machine-learning/fall_detection/scripts/inspect-data.py

@@ -14,6 +14,7 @@ def main():
         df,
         group_keys=[
             pl.col("robot_identifier"),
+            pl.col("game_phase_identifier"),
             pl.col("match_identifier"),
         ],
         features=[

tools/machine-learning/fall_detection/scripts/train.py

@@ -230,7 +230,11 @@ def train(model_type: ModelType, data_path: str) -> None:
     df = load(data_path)
     dataset = FallenDataset(
         df,
-        group_keys=[pl.col("robot_identifier"), pl.col("match_identifier")],
+        group_keys=[
+            pl.col("robot_identifier"),
+            pl.col("game_phase_identifier"),
+            pl.col("match_identifier"),
+        ],
         features=[
             pl.col(
                 "Control.main_outputs.sensor_data.inertial_measurement_unit.linear_acceleration.x"