diff --git a/docs/img/collisions.png b/docs/img/collisions.png new file mode 100644 index 0000000..8969ddb Binary files /dev/null and b/docs/img/collisions.png differ diff --git a/fignet/data_loader.py b/fignet/data_loader.py index 3388b3e..2c2367b 100644 --- a/fignet/data_loader.py +++ b/fignet/data_loader.py @@ -21,7 +21,10 @@ # SOFTWARE. +import os +import pickle from dataclasses import fields +from pathlib import Path from typing import List import numpy as np @@ -32,8 +35,6 @@ from fignet.types import EdgeType, Graph, NodeType from fignet.utils import dataclass_to_tensor, dict_to_tensor -# from torch.utils.data import DistributedSampler - def collate_fn(batch: List[Graph]): """Merge batch of graphs into one graph""" @@ -127,27 +128,46 @@ def __init__( config: dict = None, transform=None, ): + # If raw data is given, need to calculate graph connectivity on the fly + if os.path.isfile(path): + self._load_raw_data = True + + self._data = list(np.load(path, allow_pickle=True).values())[0] + self._dimension = self._data[0]["pos"].shape[-1] + self._target_length = 1 + self._input_sequence_length = input_sequence_length + + self._data_lengths = [ + x["pos"].shape[0] - input_sequence_length - self._target_length + for x in self._data + ] + self._length = sum(self._data_lengths) - self._data = list(np.load(path, allow_pickle=True).values())[0] - self._dimension = self._data[0]["pos"].shape[-1] - self._target_length = 1 - self._input_sequence_length = input_sequence_length - - self._data_lengths = [ - x["pos"].shape[0] - input_sequence_length - self._target_length - for x in self._data - ] - self._length = sum(self._data_lengths) + # pre-compute cumulative lengths + # to allow fast indexing in __getitem__ + self._precompute_cumlengths = [ + sum(self._data_lengths[:x]) + for x in range(1, len(self._data_lengths) + 1) + ] + self._precompute_cumlengths = np.array( + self._precompute_cumlengths, dtype=int + ) + # Directly load pre-calculated graphs and save time + elif os.path.isdir(path): + self._load_raw_data = False + + self._graph_ext = "pkl" + self._file_list = [ + os.path.join(path, f) + for f in os.listdir(path) + if os.path.isfile(os.path.join(path, f)) + and f.endswith(self._graph_ext) + ] + self._file_list.sort(key=lambda f: int(Path(f).stem.split("_")[1])) + self._length = len(self._file_list) + else: + raise FileNotFoundError(f"{path} not found") - # pre-compute cumulative lengths - # to allow fast indexing in __getitem__ - self._precompute_cumlengths = [ - sum(self._data_lengths[:x]) - for x in range(1, len(self._data_lengths) + 1) - ] - self._precompute_cumlengths = np.array( - self._precompute_cumlengths, dtype=int - ) self._transform = transform self._mode = mode if config is not None: @@ -164,61 +184,81 @@ def __getitem__(self, idx): elif self._mode == "trajectory": return self._get_trajectory(idx) - def _get_sample(self, idx): - """Sample one step""" - trajectory_idx = np.searchsorted( - self._precompute_cumlengths - 1, idx, side="left" - ) - # Compute index of pick along time-dimension of trajectory. - start_of_selected_trajectory = ( - self._precompute_cumlengths[trajectory_idx - 1] - if trajectory_idx != 0 - else 0 - ) - time_idx = self._input_sequence_length + ( - idx - start_of_selected_trajectory - ) - - start = time_idx - self._input_sequence_length - end = time_idx - obj_ids = dict(self._data[trajectory_idx]["obj_id"].item()) - positions = self._data[trajectory_idx]["pos"][ - start:end - ] # (seq_len, n_obj, 3) input sequence - quats = self._data[trajectory_idx]["quat"][ - start:end - ] # (seq_len, n_obj, 4) input sequence - target_posisitons = self._data[trajectory_idx]["pos"][time_idx] - target_quats = self._data[trajectory_idx]["quat"][time_idx] - poses = np.concatenate([positions, quats], axis=-1) - target_poses = np.concatenate( - [target_posisitons, target_quats], axis=-1 - ) - - scene_config = dict(self._data[trajectory_idx]["meta_data"].item()) - - connectivity_radius = self._config.get("connectivity_radius") - if connectivity_radius is not None: - scene_config.update({"connectivity_radius": connectivity_radius}) + def _load_graph(self, graph_file): + try: + with open(graph_file, "rb") as f: + sample_dict = pickle.load(f) + graph = Graph() + graph.from_dict(sample_dict) + return graph - noise_std = self._config.get("noise_std") - if noise_std is not None: - scene_config.update({"noise_std": noise_std}) + except FileNotFoundError as e: + print(e) + return None - scn = Scene(scene_config) - scn.synchronize_states( - obj_poses=poses, - obj_ids=obj_ids, - ) - graph = scn.to_graph( - target_poses=target_poses, - obj_ids=obj_ids, - noise=True, - ) + def _get_sample(self, idx): + """Sample one step""" + if self._load_raw_data: + trajectory_idx = np.searchsorted( + self._precompute_cumlengths - 1, idx, side="left" + ) + # Compute index of pick along time-dimension of trajectory. + start_of_selected_trajectory = ( + self._precompute_cumlengths[trajectory_idx - 1] + if trajectory_idx != 0 + else 0 + ) + time_idx = self._input_sequence_length + ( + idx - start_of_selected_trajectory + ) + + start = time_idx - self._input_sequence_length + end = time_idx + obj_ids = dict(self._data[trajectory_idx]["obj_id"].item()) + positions = self._data[trajectory_idx]["pos"][ + start:end + ] # (seq_len, n_obj, 3) input sequence + quats = self._data[trajectory_idx]["quat"][ + start:end + ] # (seq_len, n_obj, 4) input sequence + target_posisitons = self._data[trajectory_idx]["pos"][time_idx] + target_quats = self._data[trajectory_idx]["quat"][time_idx] + poses = np.concatenate([positions, quats], axis=-1) + target_poses = np.concatenate( + [target_posisitons, target_quats], axis=-1 + ) + + scene_config = dict(self._data[trajectory_idx]["meta_data"].item()) + + connectivity_radius = self._config.get("connectivity_radius") + if connectivity_radius is not None: + scene_config.update( + {"connectivity_radius": connectivity_radius} + ) - if self._transform is not None: - graph = self._transform(graph) - return graph + noise_std = self._config.get("noise_std") + if noise_std is not None: + scene_config.update({"noise_std": noise_std}) + + scn = Scene(scene_config) + scn.synchronize_states( + obj_poses=poses, + obj_ids=obj_ids, + ) + graph = scn.to_graph( + target_poses=target_poses, + obj_ids=obj_ids, + noise=True, + ) + + if self._transform is not None: + graph = self._transform(graph) + return graph + else: + if os.path.exists(self._file_list[idx]): + return self._load_graph(self._file_list[idx]) + else: + raise FileNotFoundError def _get_trajectory(self, idx): """Sample continuous steps""" diff --git a/fignet/types.py b/fignet/types.py index 7649525..eb11b55 100644 --- a/fignet/types.py +++ b/fignet/types.py @@ -21,33 +21,70 @@ # SOFTWARE. import enum -from dataclasses import asdict, dataclass, field +from dataclasses import asdict, dataclass, field, fields, is_dataclass from typing import Dict, Union import numpy as np +from dacite import from_dict from torch import Tensor TensorType = Union[np.ndarray, Tensor] +class MetaEnum(enum.EnumMeta): + def __contains__(cls, item): + try: + cls(item) + except ValueError: + return False + return True + + class KinematicType(enum.IntEnum): STATIC = 0 DYNAMIC = 1 SIZE = 1 -class NodeType(enum.Enum): +class NodeType(enum.Enum, metaclass=MetaEnum): MESH = "mesh" OBJECT = "object" -class EdgeType(enum.Enum): +class EdgeType(enum.Enum, metaclass=MetaEnum): MESH_MESH = "mesh-mesh" MESH_OBJ = "mesh-object" OBJ_MESH = "object-mesh" FACE_FACE = "face-face" +def key_to_string(key): + if isinstance(key, str): + return key + elif isinstance(key, enum.Enum): + return key.value + else: + raise ValueError(f"{type(key)} not supported") + + +def string_to_enum(k_str, enum_list): + for e in enum_list: + if k_str in e: + return e(k_str) + + +def to_dict(item): + d = {} + if isinstance(item, Dict): + for k, v in item.items(): + d.update({key_to_string(k): to_dict(v)}) + return d + elif is_dataclass(item): + return to_dict(asdict(item)) + else: + return item + + @dataclass class NodeFeature: position: TensorType @@ -68,4 +105,22 @@ class Graph: edge_sets: Dict[EdgeType, Edge] = field(default_factory=lambda: {}) def to_dict(self): - return asdict(self) + return to_dict(self) + + def from_dict(self, d): + for f in fields(self): + field_name = f.name + if field_name not in d: + continue + if field_name == "node_sets": + d_cls = NodeFeature + elif field_name == "edge_sets": + d_cls = Edge + for k, v in d[field_name].items(): + getattr(self, field_name).update( + { + string_to_enum(k, [NodeType, EdgeType]): from_dict( + d_cls, v + ) + } + ) diff --git a/readme.md b/readme.md index 5448727..6751c71 100644 --- a/readme.md +++ b/readme.md @@ -20,7 +20,7 @@ Graph Networks](https://arxiv.org/pdf/2212.03574)\[1\]. We try to reproduce the results from the original paper. The current version is still in an early experimental stage. -### Dependencies +### Major dependencies - [Pytorch](https://pytorch.org/) - [Pytorch3D](https://github.com/facebookresearch/pytorch3d) @@ -69,19 +69,25 @@ The dataset is stored as a .npz file. Each trajectory contains a dictionary ### Node features -For node features, we use: node velocities, +For node features, we use: 2 FD node velocities, inverse of mass, friction, restitution, object kinematic -### Calculate face connectivity +### Face-face edges We use the [hpp-fcl](https://github.com/humanoid-path-planner/hpp-fcl) library -to calculate face-face edges and features. +to detect collisions, and calculate face-face edge features based on the +detection. + +
+ +
### Graph structure and message passing Because of the object-mesh edges and the novel face-face edges, we build a -graph with two sets of nodes and three sets of edges. The message passing layer -is augmented to handle face-face messages +graph with two sets of nodes (mesh and object nodes) and four sets of edges +(mesh-mesh, mesh-object, object-mesh, face-face). The message passing layer +is augmented to handle face-face message passing. ## How to Install @@ -150,28 +156,44 @@ python scripts/generate_data.py --ep_len=100 --internal_steps=10 --total_steps=1 python scripts/generate_dataset.py --ep_len=100 --internal_steps=10 --total_steps=100000 --data_path=datasets # Generate 100k steps for testing ``` +You can pre-compute the graphs from the raw dataset beforehand so that the training runs faster +(only the training dataset). + +```bash +python scripts/preprocess_data.py --data_path=[path_to_dataset/train_dataset_name.npz] --num_workers=[default to 1] +``` + +This process takes around 8 hours with `num_workers=8`, and will create a +folder `path_to_dataset/train_dataset_name` with all +the pre-computed graphs stored inside. The dataset with 1M steps will create +960k graphs and takes around 130GB disk space. Alternatively, the pre-computed +graphs for training can also be downloaded +[here](https://cloud.dfki.de/owncloud/index.php/s/Z3ptYmdERrMmSe4). It needs to +be uncompressed. + ### 2. Run the training +For the training you need to pass in a config file; a template can be found in +[config/train.json](config/train.json). Adapt `data_path`, `test_data_path` to +the train and test dataset respectively. For train dataset, it can be the raw +dataset (npz file) or the folder containing pre-computed graphs, while the test +dataset should be a npz file. Also adapt `batch_size` and `num_workers` accordingly. + ```bash python scripts/train.py --config_file=config/train.json ``` -### 3. Render rollout +### 3. Generate animation + +The [render_model.py](scripts/render_model.py) script will sample several rollouts +with the learned simulator, and generate animation of the ground truth and +predicted trajectories. ```bash python scripts/render_model.py --model_path=[model path] --num_ep=[number of episodes] --off_screen --video_path=[video path] ``` -After training for about 300k steps (should train for more steps) with batch -size 32 (due to limited memory), we generated rollouts of 200 steps. -The rendered trajectories are shown below, with top row the ground truth and -bottom row the simulation. - -
- - - -
+:construction: **TODO:** generated gifs ## Acknowledgments diff --git a/requirements.txt b/requirements.txt index 5558ae5..f36e6a3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,3 +8,4 @@ trimesh==3.21.7 # hpp-fcl==2.4.5 numpy==1.23.1 tensorboard==2.13.0 +dacite==1.8.1 diff --git a/scripts/preprocess_data.py b/scripts/preprocess_data.py new file mode 100644 index 0000000..1e5e885 --- /dev/null +++ b/scripts/preprocess_data.py @@ -0,0 +1,102 @@ +# MIT License +# +# Copyright (c) [2024] [Zongyao Yi] +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +import argparse +import json +import os +import pickle +import sys +from pathlib import Path + +import torch +import torchvision.transforms as T +import tqdm + +from fignet.data_loader import MujocoDataset, ToTensor + +parser = argparse.ArgumentParser() +parser.add_argument("--config_file", type=str, required=True) +parser.add_argument("--data_path", type=str, required=True) +parser.add_argument("--num_workers", type=int, required=True, default=1) + +args = parser.parse_args() + +data_path = args.data_path +config_file = args.config_file +num_workers = max(args.num_workers, 1) +batch_size = min(2 * num_workers, 64) +output_path = os.path.join(Path(data_path).parent, Path(data_path).stem) +device = torch.device("cpu") + + +def collate_fn(batch): + return batch + + +def save_graph(graph, graph_i, save_path): + if isinstance(graph, list): + batch_size = len(graph) + for g_i, g in enumerate(graph): + i = graph_i * batch_size + g_i + save_graph(g, i, save_path) + else: + graph_dict = graph.to_dict() + if not os.path.exists(save_path): + os.mkdir(save_path) + file_name = os.path.join(save_path, f"graph_{graph_i}.pkl") + with open(file_name, "wb") as f: + pickle.dump(graph_dict, f) + + +if __name__ == "__main__": + try: + with open(os.path.join(os.getcwd(), args.config_file)) as f: + config = json.load(f) + except FileNotFoundError as e: + print(e) + sys.exit() + + print( + f"Parsing {data_path}. Preprocessed graphs will be stored in {output_path}" + ) + try: + dataset = MujocoDataset( + path=data_path, + mode="sample", + input_sequence_length=3, + transform=T.Compose([ToTensor(device)]), + config=config.get("data_config"), + ) + data_loader = torch.utils.data.DataLoader( + dataset, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + pin_memory=False, + collate_fn=collate_fn, + ) + for i, sample in enumerate( + tqdm.tqdm(data_loader, desc="Preprocessing data") + ): + save_graph(sample, i, output_path) + except FileNotFoundError as e: + print(e)