From de3b9298eefe63d603cfb82109c27d23a6ef9c0f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20K=C3=BClz?= <jonathan.kuelz@tum.de>
Date: Thu, 7 Mar 2024 12:32:28 +0100
Subject: [PATCH 01/12] Adds a first draft of a generic IK function that works
 with a pretrained model on goals and robot kinematics solely defined via
 torch Tensors

---
 .gitignore                                    |  1 +
 generative_graphik/utils/api.py               | 89 +++++++++++++++++++
 .../utils/dataset_generation.py               | 13 +++
 generative_graphik/utils/get_model.py         | 38 ++++++++
 generative_graphik/utils/torch_to_graphik.py  | 47 ++++++++++
 sample_config.yaml                            |  1 +
 tests/__init__.py                             | 10 +++
 tests/test_ik_api.py                          | 81 +++++++++++++++++
 8 files changed, 280 insertions(+)
 create mode 100644 generative_graphik/utils/api.py
 create mode 100644 generative_graphik/utils/get_model.py
 create mode 100644 generative_graphik/utils/torch_to_graphik.py
 create mode 100644 sample_config.yaml
 create mode 100644 tests/__init__.py
 create mode 100644 tests/test_ik_api.py

diff --git a/.gitignore b/.gitignore
index 9dcb1b5..355d02a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+config.yaml
 saved_models/
 datasets/
 **/results/
diff --git a/generative_graphik/utils/api.py b/generative_graphik/utils/api.py
new file mode 100644
index 0000000..2857c3f
--- /dev/null
+++ b/generative_graphik/utils/api.py
@@ -0,0 +1,89 @@
+import itertools
+from typing import Callable
+
+from liegroups.numpy.se3 import SE3Matrix
+import torch
+
+from graphik.graphs import ProblemGraphRevolute
+from graphik.robots import RobotRevolute
+from graphik.utils import graph_from_pos
+from generative_graphik.utils.dataset_generation import generate_data_point_from_pose
+from generative_graphik.utils.get_model import get_model
+from generative_graphik.utils.torch_to_graphik import joint_transforms_to_t_zero
+
+
+def _default_cost_function(T_desired: torch.Tensor, T_eef: torch.Tensor) -> torch.Tensor:
+    """
+    The default cost function for the inverse kinematics problem. It is the sum of the squared errors between the
+    desired and actual end-effector poses.
+
+    :param T_desired: The desired end-effector pose.
+    :param T_eef: The actual end-effector pose.
+    :return: The cost.
+    """
+    return torch.sum((T_desired - T_eef) ** 2)
+
+
+def ik(kinematic_chains: torch.tensor,
+       goals: torch.tensor,
+       samples: int = 16,
+       return_all: bool = False,
+       ik_cost_function: Callable = _default_cost_function) -> torch.Tensor:
+    """
+    This function takes robot kinematics and any number of goals and solves the inverse kinematics, using graphIK.
+
+    :param kinematic_chains: A tensor of shape (nR, N, 4, 4) containing the joint transformations of nR robots with N
+        joints each.
+    :param goals: A tensor of shape (nG, 4, 4) containing the desired end-effector poses.
+    :param samples: The number of samples to use for the forward pass of the model.
+    :param return_all: If True, returns all the samples from the forward pass, so the resulting tensor has a shape
+        nR x nG x samples x nJ. If False, returns the best one only, so the resulting tensor has a shape nR x nG x nJ.
+    :param ik_cost_function: The cost function to use for the inverse kinematics problem if return_all is False.
+    """
+    device = kinematic_chains.device
+    model = get_model().to(device)
+
+    assert len(kinematic_chains.shape) == 4, f'Expected 4D tensor, got {kinematic_chains.shape}'
+    nr, nj, _, _ = kinematic_chains.shape
+
+    t_zeros = {i: joint_transforms_to_t_zero(kinematic_chains[i], [f'p{j}' for j in range(1 + nj)]) for i in range(nr)}
+    robots = {i: RobotRevolute({'num_joints': nj, 'T_zero': t_zeros[i]}) for i in range(nr)}
+    graphs = {i: ProblemGraphRevolute(robots[i]) for i in range(nr)}
+    data = {i: {j: generate_data_point_from_pose(graphs[i], goals[j]) for j in range(len(goals))} for i in range(nr)}
+    if return_all:
+        q = torch.zeros((nr, len(goals), samples, nj), device=device)
+    else:
+        q = torch.zeros((nr, len(goals), nj), device=device)
+
+    for i, j in itertools.product(range(nr), range(len(goals))):
+        graph = graphs[i]
+        robot = robots[i]
+        problem = data[i][j]
+        problem = model.preprocess(problem)
+        P_all = (
+            model.forward_eval(
+                x=problem.pos,
+                h=torch.cat((problem.type, problem.goal_data_repeated_per_node), dim=-1),
+                edge_attr=problem.edge_attr,
+                edge_attr_partial=problem.edge_attr_partial,
+                edge_index=problem.edge_index_full,
+                partial_goal_mask=problem.partial_goal_mask,
+                nodes_per_single_graph=int(problem.num_nodes / 1),
+                batch_size=1,
+                num_samples=samples
+            )
+        ).cpu().detach().numpy()
+        best = float('inf')
+        for k, p_k in enumerate(P_all):
+            q_k = graph.joint_variables(graph_from_pos(p_k, graph.node_ids),
+                                        {robot.end_effectors[0]: SE3Matrix.from_matrix(goals[j].detach().cpu().numpy(), normalize=True)})
+            q_k = torch.tensor([q_k[key] for key in robot.joint_ids[1:]], device=device)
+            if return_all:
+                q[i, j, k] = torch.tensor(q_k, device=device)
+                continue
+            T_ee = graph.robot.pose(q_k, robot.end_effectors[-1])
+            cost = ik_cost_function(goals[j], T_ee)
+            if cost < best:
+                best = cost
+                q[i, j] = torch.tensor(q_k, device=device)
+    return q
diff --git a/generative_graphik/utils/dataset_generation.py b/generative_graphik/utils/dataset_generation.py
index 8ab805d..b4b79ae 100644
--- a/generative_graphik/utils/dataset_generation.py
+++ b/generative_graphik/utils/dataset_generation.py
@@ -1,4 +1,7 @@
 from typing import List, Union
+
+from liegroups.numpy.se2 import SE2Matrix
+from liegroups.numpy.se3 import SE3Matrix
 import numpy as np
 import os
 from tqdm import tqdm
@@ -58,6 +61,16 @@ class StructData:
     T0: Union[List[torch.Tensor], torch.Tensor]
 
 def generate_data_point_from_pose(graph, T_ee):
+    """
+    Generates a data point (~problem) from a problem graph and a desired end-effector pose.
+    """
+    if isinstance(T_ee, torch.Tensor):
+        T_ee = T_ee.detach().cpu().numpy()
+    if isinstance(T_ee, np.ndarray):
+        if T_ee.shape == (4, 4):
+            T_ee = SE3Matrix.from_matrix(T_ee, normalize=True)
+        else:
+            raise ValueError(f"Expected T_ee to be of shape (4, 4) or be SEMatrix, got {T_ee.shape}")
     struct_data = generate_struct_data(graph)
     num_joints = torch.tensor([struct_data.num_joints])
     edge_index_full = struct_data.edge_index_full
diff --git a/generative_graphik/utils/get_model.py b/generative_graphik/utils/get_model.py
new file mode 100644
index 0000000..4e16304
--- /dev/null
+++ b/generative_graphik/utils/get_model.py
@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+# Author: Jonathan Külz
+# Date: 06.03.24
+from argparse import Namespace
+import json
+from pathlib import Path
+from typing import Dict
+
+import torch
+import yaml
+
+from generative_graphik.model import Model
+
+_model = None  # Use get_model to access the model
+PROJECT_DIR = Path(__file__).resolve().parents[2]
+CONFIG_DIR = PROJECT_DIR.joinpath('config.yaml')
+
+
+def get_config() -> Dict:
+    """Loads the configuration file"""
+    with CONFIG_DIR.open('r') as f:
+        return yaml.safe_load(f)
+
+
+def get_model() -> Model:
+    """Loads the model specified in the configuration file or returns the cached model."""
+    global _model
+    if _model is not None:
+        return _model
+    config = get_config()
+    d = Path(config['model'])
+    state_dict = torch.load(d.joinpath('net.pth'))
+    with d.joinpath('hyperparameters.txt').open('r') as f:
+        args = Namespace(**json.load(f))
+    model = Model(args)
+    model.load_state_dict(state_dict)
+    _model = model
+    return model
diff --git a/generative_graphik/utils/torch_to_graphik.py b/generative_graphik/utils/torch_to_graphik.py
new file mode 100644
index 0000000..520831c
--- /dev/null
+++ b/generative_graphik/utils/torch_to_graphik.py
@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+# Author: Jonathan Külz
+# Date: 06.03.24
+import itertools
+from typing import Dict, Sequence
+
+from liegroups.numpy.se3 import SE3Matrix
+import torch
+from torch_geometric.data import Data
+
+from generative_graphik.utils.dataset_generation import StructData
+
+
+def define_ik_data(robot_data: StructData, goals: torch.Tensor) -> Data:
+    """
+    This function takes a robot and a set of goals and returns a data point for every goal.
+
+    :param robot_data: A StructData object containing the robot's kinematics.
+    :param goals: A tensor of shape (nG, 4, 4) containing the desired end-effector poses.
+    """
+    pass
+
+
+def joint_transforms_from_t_zeros(T_zero: Dict[str, SE3Matrix], keys: Sequence[str]) -> torch.Tensor:
+    """Assumes that joints are alphabetically sorted"""
+    ret = torch.zeros((len(T_zero) - 1, 4, 4))
+    for i in range(1, len(keys)):
+        ret[i - 1] = torch.Tensor(T_zero[keys[i-1]].inv().dot(T_zero[keys[i]]).as_matrix())
+    return ret
+
+
+def joint_transforms_to_t_zero(transforms: torch.Tensor, keys: Sequence[str]) -> Dict[str, SE3Matrix]:
+    """
+    This function takes a tensor of joint transformations and returns the t_zero tensor, which describes the joint
+    pose in the world frame for the zero configuration.
+
+    :param transforms: A tensor of shape (nJ, 4, 4).
+    :param keys: The keys to use for the joint names. Assumes the first key is for the world frame, thus it will be
+        set to the identity.
+    """
+    nj = transforms.shape[0]
+    t_zero = transforms.clone()
+    for i in range(1, nj):
+        t_zero[i] = t_zero[i - 1] @ t_zero[i]
+    t_zero = {keys[i+1]: SE3Matrix.from_matrix(t_zero[i].detach().cpu().numpy()) for i in range(nj)}
+    t_zero[keys[0]] = SE3Matrix.identity()
+    return t_zero
diff --git a/sample_config.yaml b/sample_config.yaml
new file mode 100644
index 0000000..b2196d4
--- /dev/null
+++ b/sample_config.yaml
@@ -0,0 +1 @@
+model: '<path_to_your_pretrained_model>'
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..c7f49b8
--- /dev/null
+++ b/tests/__init__.py
@@ -0,0 +1,10 @@
+#!/usr/bin/env python3
+# Author: Jonathan Külz
+# Date: 07.03.24
+
+def main():
+    pass
+
+
+if __name__ == '__main__':
+    main()
diff --git a/tests/test_ik_api.py b/tests/test_ik_api.py
new file mode 100644
index 0000000..325f656
--- /dev/null
+++ b/tests/test_ik_api.py
@@ -0,0 +1,81 @@
+import itertools
+from typing import Tuple
+import unittest
+
+import numpy as np
+import torch
+
+from generative_graphik.utils.dataset_generation import random_revolute_robot_graph
+from generative_graphik.utils.api import ik
+from generative_graphik.utils.get_model import get_model
+from generative_graphik.utils.torch_to_graphik import joint_transforms_from_t_zeros, joint_transforms_to_t_zero
+
+
+class ApiTests(unittest.TestCase):
+
+    def setUp(self):
+        try:
+            self.model = get_model()
+        except Exception as e:
+            print(e)
+            raise e
+
+    @staticmethod
+    def ik_error(T_desired: torch.Tensor, T_eef: torch.Tensor) -> Tuple[float, float]:
+        r_desired = T_desired[:3, :3]
+        r_eef = T_eef[:3, :3]
+        t_desired = T_desired[:3, 3]
+        t_eef = T_eef[:3, 3]
+        e_rot = torch.arccos((torch.trace(r_desired.T @ r_eef) - 1) / 2) * 180 / np.pi  # degrees
+        e_trans = torch.norm(t_desired - t_eef)  # meters
+        return e_rot.item(), e_trans.item()
+
+    def test_conversions(self, N=10, dof=6):
+        for _ in range(N):
+            g = random_revolute_robot_graph(dof)
+            T_zero = g.robot.from_dh_params(g.robot.params)
+            transforms = joint_transforms_from_t_zeros(T_zero, keys=g.robot.joint_ids)
+            T_zero_reconstructed = joint_transforms_to_t_zero(transforms, keys=g.robot.joint_ids)
+            for key in T_zero:
+                self.assertTrue(np.allclose(T_zero[key].as_matrix(), T_zero_reconstructed[key].as_matrix()))
+
+    def test_ik_api(self, nR: int = 8, nG: int = 8, samples: int = 8, dof: int = 6):
+        graphs = [random_revolute_robot_graph(dof) for _ in range(nR)]
+        goals = dict.fromkeys(range(nR), None)
+        for i, j in itertools.product(range(nR), range(nG)):
+            if j == 0:
+                goals[i] = []
+            q = torch.rand(dof + 1) * 2 * torch.pi - torch.pi
+            # q[-1] = 0
+            angles = {jnt: q_jnt.item() for jnt, q_jnt in zip(graphs[i].robot.joint_ids, q)}
+            T = graphs[i].robot.pose(angles, graphs[i].robot.end_effectors[-1])
+            goals[i].append(torch.Tensor(T.as_matrix()))
+
+        for i, g in enumerate(graphs):
+            T_zero_native = g.robot.from_dh_params(g.robot.params)
+            transforms = joint_transforms_from_t_zeros(T_zero_native, keys=g.robot.joint_ids)
+            transforms = torch.unsqueeze(transforms, 0)  # T=D=
+            sol = ik(transforms, torch.stack(goals[i]), samples=samples, return_all=True)
+
+            trans_errors = list()
+            rot_errors = list()
+
+            for k, l in itertools.product(range(nG), range(samples)):
+                q_kl = {jnt: sol[0, k, l, m].item() for m, jnt in enumerate(g.robot.joint_ids[1:])}
+                q_kl['p0'] = 0
+                T = g.robot.pose(q_kl, g.robot.end_effectors[-1]).as_matrix()
+                e_rot, e_trans = self.ik_error(goals[i][k], torch.Tensor(T))
+                rot_errors.append(e_rot)
+                trans_errors.append(e_trans)
+
+            # Get at least one good solution
+            self.assertLessEqual(np.min(rot_errors), 2)
+            self.assertLessEqual(np.min(trans_errors), 0.05)
+
+            # Is it significantly better than average?
+            self.assertLessEqual(np.mean(rot_errors), 45)
+            self.assertLessEqual(np.min(trans_errors), np.mean([np.linalg.norm(goals[i][j][:3, 3]) for j in range(nG)]) / 10)
+
+
+if __name__ == '__main__':
+    unittest.main()

From 40efbad506ea2b8bc1000e0d02aa96398dde7199 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20K=C3=BClz?= <jonathan.kuelz@tum.de>
Date: Thu, 7 Mar 2024 12:40:51 +0100
Subject: [PATCH 02/12] removes leftover headers

---
 generative_graphik/utils/get_model.py        | 3 ---
 generative_graphik/utils/torch_to_graphik.py | 4 ----
 tests/__init__.py                            | 4 ----
 3 files changed, 11 deletions(-)

diff --git a/generative_graphik/utils/get_model.py b/generative_graphik/utils/get_model.py
index 4e16304..bf5b777 100644
--- a/generative_graphik/utils/get_model.py
+++ b/generative_graphik/utils/get_model.py
@@ -1,6 +1,3 @@
-#!/usr/bin/env python3
-# Author: Jonathan Külz
-# Date: 06.03.24
 from argparse import Namespace
 import json
 from pathlib import Path
diff --git a/generative_graphik/utils/torch_to_graphik.py b/generative_graphik/utils/torch_to_graphik.py
index 520831c..8d013f6 100644
--- a/generative_graphik/utils/torch_to_graphik.py
+++ b/generative_graphik/utils/torch_to_graphik.py
@@ -1,7 +1,3 @@
-#!/usr/bin/env python3
-# Author: Jonathan Külz
-# Date: 06.03.24
-import itertools
 from typing import Dict, Sequence
 
 from liegroups.numpy.se3 import SE3Matrix
diff --git a/tests/__init__.py b/tests/__init__.py
index c7f49b8..a53c62e 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -1,7 +1,3 @@
-#!/usr/bin/env python3
-# Author: Jonathan Külz
-# Date: 07.03.24
-
 def main():
     pass
 

From f4d341a346f27ce49d4662465e70ca39f7790c3b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20K=C3=BClz?= <jonathan.kuelz@tum.de>
Date: Thu, 7 Mar 2024 12:47:49 +0100
Subject: [PATCH 03/12] Adds proper exception when the config for get_model is
 not yet set up

---
 tests/test_ik_api.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/tests/test_ik_api.py b/tests/test_ik_api.py
index 325f656..adde1ef 100644
--- a/tests/test_ik_api.py
+++ b/tests/test_ik_api.py
@@ -12,16 +12,25 @@
 
 
 class ApiTests(unittest.TestCase):
+    """
+    Tests the generative_graphik.utils.api functionalities.
+    """
 
     def setUp(self):
         try:
             self.model = get_model()
-        except Exception as e:
-            print(e)
-            raise e
+        except FileNotFoundError as exe:
+            print(exe)
+            if exe.filename.split('/')[-1] == 'config.yaml':
+                raise FileNotFoundError("No configuration file found. Create a config.yaml file similar to "
+                                        "sample_config.yaml and place it in the root of the project.")
+            else:
+                raise FileNotFoundError("No model found. Train a model and place it in the directory specified in the "
+                                        "config.yaml file.")
 
     @staticmethod
     def ik_error(T_desired: torch.Tensor, T_eef: torch.Tensor) -> Tuple[float, float]:
+        """Utility function to compute the error between two SE3 matrices."""
         r_desired = T_desired[:3, :3]
         r_eef = T_eef[:3, :3]
         t_desired = T_desired[:3, 3]
@@ -31,6 +40,7 @@ def ik_error(T_desired: torch.Tensor, T_eef: torch.Tensor) -> Tuple[float, float
         return e_rot.item(), e_trans.item()
 
     def test_conversions(self, N=10, dof=6):
+        """Test that the joint transforms (torch) can be converted to T_zero (dict of SE3) and back."""
         for _ in range(N):
             g = random_revolute_robot_graph(dof)
             T_zero = g.robot.from_dh_params(g.robot.params)
@@ -40,6 +50,10 @@ def test_conversions(self, N=10, dof=6):
                 self.assertTrue(np.allclose(T_zero[key].as_matrix(), T_zero_reconstructed[key].as_matrix()))
 
     def test_ik_api(self, nR: int = 8, nG: int = 8, samples: int = 8, dof: int = 6):
+        """
+        Test the inverse kinematics API, i.e., an inverse kinematics functionality that is framework-agnostic and does
+        not require the user to know the details of the generative_graphik approach.
+        """
         graphs = [random_revolute_robot_graph(dof) for _ in range(nR)]
         goals = dict.fromkeys(range(nR), None)
         for i, j in itertools.product(range(nR), range(nG)):
@@ -72,7 +86,7 @@ def test_ik_api(self, nR: int = 8, nG: int = 8, samples: int = 8, dof: int = 6):
             self.assertLessEqual(np.min(rot_errors), 2)
             self.assertLessEqual(np.min(trans_errors), 0.05)
 
-            # Is it significantly better than average?
+            # Is it significantly better than random? (educated guess of what a random precision would be)
             self.assertLessEqual(np.mean(rot_errors), 45)
             self.assertLessEqual(np.min(trans_errors), np.mean([np.linalg.norm(goals[i][j][:3, 3]) for j in range(nG)]) / 10)
 

From a989c72406b5d8591c50cce780a6a3d6088861bf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20K=C3=BClz?= <jonathan.kuelz@tum.de>
Date: Thu, 7 Mar 2024 14:08:41 +0100
Subject: [PATCH 04/12] move test tensors to GPU

---
 generative_graphik/utils/api.py              |  2 +-
 generative_graphik/utils/torch_to_graphik.py |  6 +--
 sample_config.yaml                           |  2 +-
 tests/test_ik_api.py                         | 46 +++++++++++++++-----
 4 files changed, 40 insertions(+), 16 deletions(-)

diff --git a/generative_graphik/utils/api.py b/generative_graphik/utils/api.py
index 2857c3f..827e0f9 100644
--- a/generative_graphik/utils/api.py
+++ b/generative_graphik/utils/api.py
@@ -49,7 +49,7 @@ def ik(kinematic_chains: torch.tensor,
     t_zeros = {i: joint_transforms_to_t_zero(kinematic_chains[i], [f'p{j}' for j in range(1 + nj)]) for i in range(nr)}
     robots = {i: RobotRevolute({'num_joints': nj, 'T_zero': t_zeros[i]}) for i in range(nr)}
     graphs = {i: ProblemGraphRevolute(robots[i]) for i in range(nr)}
-    data = {i: {j: generate_data_point_from_pose(graphs[i], goals[j]) for j in range(len(goals))} for i in range(nr)}
+    data = {i: {j: generate_data_point_from_pose(graphs[i], goals[j]).to(device) for j in range(len(goals))} for i in range(nr)}
     if return_all:
         q = torch.zeros((nr, len(goals), samples, nj), device=device)
     else:
diff --git a/generative_graphik/utils/torch_to_graphik.py b/generative_graphik/utils/torch_to_graphik.py
index 8d013f6..3993897 100644
--- a/generative_graphik/utils/torch_to_graphik.py
+++ b/generative_graphik/utils/torch_to_graphik.py
@@ -17,11 +17,11 @@ def define_ik_data(robot_data: StructData, goals: torch.Tensor) -> Data:
     pass
 
 
-def joint_transforms_from_t_zeros(T_zero: Dict[str, SE3Matrix], keys: Sequence[str]) -> torch.Tensor:
+def joint_transforms_from_t_zeros(T_zero: Dict[str, SE3Matrix], keys: Sequence[str], device: str = None) -> torch.Tensor:
     """Assumes that joints are alphabetically sorted"""
-    ret = torch.zeros((len(T_zero) - 1, 4, 4))
+    ret = torch.zeros((len(T_zero) - 1, 4, 4), device=device)
     for i in range(1, len(keys)):
-        ret[i - 1] = torch.Tensor(T_zero[keys[i-1]].inv().dot(T_zero[keys[i]]).as_matrix())
+        ret[i - 1] = torch.tensor(T_zero[keys[i-1]].inv().dot(T_zero[keys[i]]).as_matrix(), device=device)
     return ret
 
 
diff --git a/sample_config.yaml b/sample_config.yaml
index b2196d4..8f1cfcf 100644
--- a/sample_config.yaml
+++ b/sample_config.yaml
@@ -1 +1 @@
-model: '<path_to_your_pretrained_model>'
+model: '<path_to_your_pretrained_model_directory>'
diff --git a/tests/test_ik_api.py b/tests/test_ik_api.py
index adde1ef..d64bb97 100644
--- a/tests/test_ik_api.py
+++ b/tests/test_ik_api.py
@@ -1,4 +1,5 @@
 import itertools
+from time import time
 from typing import Tuple
 import unittest
 
@@ -17,8 +18,13 @@ class ApiTests(unittest.TestCase):
     """
 
     def setUp(self):
+        """
+        These tests rely on a trained model that needs to be present.
+        Its location should be specified in the config.yaml file.
+        """
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         try:
-            self.model = get_model()
+            self.model = get_model().to(self.device)
         except FileNotFoundError as exe:
             print(exe)
             if exe.filename.split('/')[-1] == 'config.yaml':
@@ -44,33 +50,38 @@ def test_conversions(self, N=10, dof=6):
         for _ in range(N):
             g = random_revolute_robot_graph(dof)
             T_zero = g.robot.from_dh_params(g.robot.params)
-            transforms = joint_transforms_from_t_zeros(T_zero, keys=g.robot.joint_ids)
+            transforms = joint_transforms_from_t_zeros(T_zero, keys=g.robot.joint_ids, device=self.device)
             T_zero_reconstructed = joint_transforms_to_t_zero(transforms, keys=g.robot.joint_ids)
             for key in T_zero:
                 self.assertTrue(np.allclose(T_zero[key].as_matrix(), T_zero_reconstructed[key].as_matrix()))
 
-    def test_ik_api(self, nR: int = 8, nG: int = 8, samples: int = 8, dof: int = 6):
+    def test_ik_api(self, nR: int = 16, nG: int = 16, samples: int = 16, dof: int = 6):
         """
         Test the inverse kinematics API, i.e., an inverse kinematics functionality that is framework-agnostic and does
         not require the user to know the details of the generative_graphik approach.
         """
+        tic = time()
+        t_eval = 0
         graphs = [random_revolute_robot_graph(dof) for _ in range(nR)]
         goals = dict.fromkeys(range(nR), None)
         for i, j in itertools.product(range(nR), range(nG)):
             if j == 0:
                 goals[i] = []
             q = torch.rand(dof + 1) * 2 * torch.pi - torch.pi
-            # q[-1] = 0
             angles = {jnt: q_jnt.item() for jnt, q_jnt in zip(graphs[i].robot.joint_ids, q)}
             T = graphs[i].robot.pose(angles, graphs[i].robot.end_effectors[-1])
-            goals[i].append(torch.Tensor(T.as_matrix()))
+            goals[i].append(torch.Tensor(T.as_matrix()).to(self.device))
 
+        all_cost = list()
+        best_cost = list()
         for i, g in enumerate(graphs):
             T_zero_native = g.robot.from_dh_params(g.robot.params)
-            transforms = joint_transforms_from_t_zeros(T_zero_native, keys=g.robot.joint_ids)
-            transforms = torch.unsqueeze(transforms, 0)  # T=D=
+            transforms = joint_transforms_from_t_zeros(T_zero_native, keys=g.robot.joint_ids, device=self.device)
+            transforms = torch.unsqueeze(transforms, 0)
             sol = ik(transforms, torch.stack(goals[i]), samples=samples, return_all=True)
 
+            t_eval_start = time()
+
             trans_errors = list()
             rot_errors = list()
 
@@ -78,17 +89,30 @@ def test_ik_api(self, nR: int = 8, nG: int = 8, samples: int = 8, dof: int = 6):
                 q_kl = {jnt: sol[0, k, l, m].item() for m, jnt in enumerate(g.robot.joint_ids[1:])}
                 q_kl['p0'] = 0
                 T = g.robot.pose(q_kl, g.robot.end_effectors[-1]).as_matrix()
-                e_rot, e_trans = self.ik_error(goals[i][k], torch.Tensor(T))
+                e_rot, e_trans = self.ik_error(goals[i][k], torch.tensor(T, device=self.device, dtype=torch.float32))
                 rot_errors.append(e_rot)
                 trans_errors.append(e_trans)
 
+            cost = [trans + rot * 0.05 / 2 for trans, rot in zip(trans_errors, rot_errors)]  # 2° ~ 0.05m
             # Get at least one good solution
-            self.assertLessEqual(np.min(rot_errors), 2)
-            self.assertLessEqual(np.min(trans_errors), 0.05)
+            self.assertLessEqual(np.min(rot_errors), 1)
+            self.assertLessEqual(np.min(trans_errors), 0.03)
+            self.assertLessEqual(np.min(cost), 0.05)
 
             # Is it significantly better than random? (educated guess of what a random precision would be)
             self.assertLessEqual(np.mean(rot_errors), 45)
-            self.assertLessEqual(np.min(trans_errors), np.mean([np.linalg.norm(goals[i][j][:3, 3]) for j in range(nG)]) / 10)
+            self.assertLessEqual(np.min(trans_errors),
+                                 np.mean([np.linalg.norm(goals[i][j][:3, 3].detach().cpu().numpy()) for j in range(nG)]) / 10)
+            best_cost.append(np.min(cost))
+            all_cost.append(np.mean(cost))
+            t_eval += time() - t_eval_start
+
+        toc = time()
+        delta_t = toc - tic - t_eval
+        print(f"\n\nIK Test took {delta_t:.2f} seconds. That's {1000 * delta_t / (nR * nG):.2f} ms per goal "
+              f"or {1000 * delta_t / nR:.2f} ms per robot.")
+        print(f"Mean best: {np.mean(best_cost):.3f}, Std: {np.std(best_cost):.3f}")
+        print(f"Mean cost: {np.mean(all_cost):.3f}, Std: {np.std(all_cost):.3f}")
 
 
 if __name__ == '__main__':

From bac332e77b706372369d896956ebbe503d1cc4a0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20K=C3=BClz?= <jonathan.kuelz@tum.de>
Date: Thu, 7 Mar 2024 14:48:45 +0100
Subject: [PATCH 05/12] Can we move more to torch? Starting to try

---
 generative_graphik/utils/api.py              |  9 +++++----
 generative_graphik/utils/torch_to_graphik.py | 16 ++++++++++++----
 tests/test_ik_api.py                         |  4 ++--
 3 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/generative_graphik/utils/api.py b/generative_graphik/utils/api.py
index 827e0f9..067e372 100644
--- a/generative_graphik/utils/api.py
+++ b/generative_graphik/utils/api.py
@@ -46,7 +46,7 @@ def ik(kinematic_chains: torch.tensor,
     assert len(kinematic_chains.shape) == 4, f'Expected 4D tensor, got {kinematic_chains.shape}'
     nr, nj, _, _ = kinematic_chains.shape
 
-    t_zeros = {i: joint_transforms_to_t_zero(kinematic_chains[i], [f'p{j}' for j in range(1 + nj)]) for i in range(nr)}
+    t_zeros = {i: joint_transforms_to_t_zero(kinematic_chains[i], [f'p{j}' for j in range(1 + nj)], se3type='numpy') for i in range(nr)}
     robots = {i: RobotRevolute({'num_joints': nj, 'T_zero': t_zeros[i]}) for i in range(nr)}
     graphs = {i: ProblemGraphRevolute(robots[i]) for i in range(nr)}
     data = {i: {j: generate_data_point_from_pose(graphs[i], goals[j]).to(device) for j in range(len(goals))} for i in range(nr)}
@@ -72,11 +72,12 @@ def ik(kinematic_chains: torch.tensor,
                 batch_size=1,
                 num_samples=samples
             )
-        ).cpu().detach().numpy()
+        )
         best = float('inf')
         for k, p_k in enumerate(P_all):
-            q_k = graph.joint_variables(graph_from_pos(p_k, graph.node_ids),
-                                        {robot.end_effectors[0]: SE3Matrix.from_matrix(goals[j].detach().cpu().numpy(), normalize=True)})
+            q_k = graph.joint_variables(graph_from_pos(p_k.detach().cpu().numpy(), graph.node_ids),
+                                        {robot.end_effectors[0]: SE3Matrix.from_matrix(goals[j].detach().cpu().numpy(),
+                                                                                       normalize=True)})
             q_k = torch.tensor([q_k[key] for key in robot.joint_ids[1:]], device=device)
             if return_all:
                 q[i, j, k] = torch.tensor(q_k, device=device)
diff --git a/generative_graphik/utils/torch_to_graphik.py b/generative_graphik/utils/torch_to_graphik.py
index 3993897..5ad4742 100644
--- a/generative_graphik/utils/torch_to_graphik.py
+++ b/generative_graphik/utils/torch_to_graphik.py
@@ -1,6 +1,7 @@
-from typing import Dict, Sequence
+from typing import Dict, Sequence, Union
 
 from liegroups.numpy.se3 import SE3Matrix
+from liegroups.torch.se3 import SE3Matrix as SE3MatrixTorch
 import torch
 from torch_geometric.data import Data
 
@@ -25,7 +26,9 @@ def joint_transforms_from_t_zeros(T_zero: Dict[str, SE3Matrix], keys: Sequence[s
     return ret
 
 
-def joint_transforms_to_t_zero(transforms: torch.Tensor, keys: Sequence[str]) -> Dict[str, SE3Matrix]:
+def joint_transforms_to_t_zero(transforms: torch.Tensor,
+                               keys: Sequence[str],
+                               se3type: str = 'numpy') -> Dict[str, Union[SE3Matrix, SE3MatrixTorch]]:
     """
     This function takes a tensor of joint transformations and returns the t_zero tensor, which describes the joint
     pose in the world frame for the zero configuration.
@@ -33,11 +36,16 @@ def joint_transforms_to_t_zero(transforms: torch.Tensor, keys: Sequence[str]) ->
     :param transforms: A tensor of shape (nJ, 4, 4).
     :param keys: The keys to use for the joint names. Assumes the first key is for the world frame, thus it will be
         set to the identity.
+    :param se3type: The type of SE3 matrix to use. Either 'numpy' or 'torch'.
     """
     nj = transforms.shape[0]
     t_zero = transforms.clone()
     for i in range(1, nj):
         t_zero[i] = t_zero[i - 1] @ t_zero[i]
-    t_zero = {keys[i+1]: SE3Matrix.from_matrix(t_zero[i].detach().cpu().numpy()) for i in range(nj)}
-    t_zero[keys[0]] = SE3Matrix.identity()
+    if se3type == 'torch':
+        t_zero = {keys[i+1]: SE3MatrixTorch.from_matrix(t_zero[i]) for i in range(nj)}
+        t_zero[keys[0]] = SE3MatrixTorch.identity()
+    else:
+        t_zero = {keys[i+1]: SE3Matrix.from_matrix(t_zero[i].detach().cpu().numpy()) for i in range(nj)}
+        t_zero[keys[0]] = SE3Matrix.identity()
     return t_zero
diff --git a/tests/test_ik_api.py b/tests/test_ik_api.py
index d64bb97..d78cd95 100644
--- a/tests/test_ik_api.py
+++ b/tests/test_ik_api.py
@@ -55,7 +55,7 @@ def test_conversions(self, N=10, dof=6):
             for key in T_zero:
                 self.assertTrue(np.allclose(T_zero[key].as_matrix(), T_zero_reconstructed[key].as_matrix()))
 
-    def test_ik_api(self, nR: int = 16, nG: int = 16, samples: int = 16, dof: int = 6):
+    def test_ik_api(self, nR: int = 32, nG: int = 16, samples: int = 32, dof: int = 6):
         """
         Test the inverse kinematics API, i.e., an inverse kinematics functionality that is framework-agnostic and does
         not require the user to know the details of the generative_graphik approach.
@@ -111,7 +111,7 @@ def test_ik_api(self, nR: int = 16, nG: int = 16, samples: int = 16, dof: int =
         delta_t = toc - tic - t_eval
         print(f"\n\nIK Test took {delta_t:.2f} seconds. That's {1000 * delta_t / (nR * nG):.2f} ms per goal "
               f"or {1000 * delta_t / nR:.2f} ms per robot.")
-        print(f"Mean best: {np.mean(best_cost):.3f}, Std: {np.std(best_cost):.3f}")
+        print(f"Mean best: {np.mean(best_cost):.4f}, Std: {np.std(best_cost):.4f}")
         print(f"Mean cost: {np.mean(all_cost):.3f}, Std: {np.std(all_cost):.3f}")
 
 

From 08b6a80cf5800972d016d2f519c9b824db1c90ed Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20K=C3=BClz?= <jonathan.kuelz@tum.de>
Date: Thu, 7 Mar 2024 15:36:33 +0100
Subject: [PATCH 06/12] Change the function signature: Provide different goals
 for different robots

---
 generative_graphik/utils/api.py | 22 +++++++++++-----------
 tests/test_ik_api.py            |  2 +-
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/generative_graphik/utils/api.py b/generative_graphik/utils/api.py
index 067e372..f4e1dc0 100644
--- a/generative_graphik/utils/api.py
+++ b/generative_graphik/utils/api.py
@@ -34,7 +34,7 @@ def ik(kinematic_chains: torch.tensor,
 
     :param kinematic_chains: A tensor of shape (nR, N, 4, 4) containing the joint transformations of nR robots with N
         joints each.
-    :param goals: A tensor of shape (nG, 4, 4) containing the desired end-effector poses.
+    :param goals: A tensor of shape (nR, nG, 4, 4) containing the desired end-effector poses.
     :param samples: The number of samples to use for the forward pass of the model.
     :param return_all: If True, returns all the samples from the forward pass, so the resulting tensor has a shape
         nR x nG x samples x nJ. If False, returns the best one only, so the resulting tensor has a shape nR x nG x nJ.
@@ -45,20 +45,21 @@ def ik(kinematic_chains: torch.tensor,
 
     assert len(kinematic_chains.shape) == 4, f'Expected 4D tensor, got {kinematic_chains.shape}'
     nr, nj, _, _ = kinematic_chains.shape
+    _, nG, _, _ = goals.shape
 
     t_zeros = {i: joint_transforms_to_t_zero(kinematic_chains[i], [f'p{j}' for j in range(1 + nj)], se3type='numpy') for i in range(nr)}
     robots = {i: RobotRevolute({'num_joints': nj, 'T_zero': t_zeros[i]}) for i in range(nr)}
     graphs = {i: ProblemGraphRevolute(robots[i]) for i in range(nr)}
-    data = {i: {j: generate_data_point_from_pose(graphs[i], goals[j]).to(device) for j in range(len(goals))} for i in range(nr)}
     if return_all:
-        q = torch.zeros((nr, len(goals), samples, nj), device=device)
+        q = torch.zeros((nr, nG, samples, nj), device=device)
     else:
-        q = torch.zeros((nr, len(goals), nj), device=device)
+        q = torch.zeros((nr, nG, nj), device=device)
 
-    for i, j in itertools.product(range(nr), range(len(goals))):
+    for i, j in itertools.product(range(nr), range(nG)):
         graph = graphs[i]
         robot = robots[i]
-        problem = data[i][j]
+        goal = goals[i, j]
+        problem = generate_data_point_from_pose(graph, goal).to(device)
         problem = model.preprocess(problem)
         P_all = (
             model.forward_eval(
@@ -76,15 +77,14 @@ def ik(kinematic_chains: torch.tensor,
         best = float('inf')
         for k, p_k in enumerate(P_all):
             q_k = graph.joint_variables(graph_from_pos(p_k.detach().cpu().numpy(), graph.node_ids),
-                                        {robot.end_effectors[0]: SE3Matrix.from_matrix(goals[j].detach().cpu().numpy(),
+                                        {robot.end_effectors[0]: SE3Matrix.from_matrix(goal.detach().cpu().numpy(),
                                                                                        normalize=True)})
-            q_k = torch.tensor([q_k[key] for key in robot.joint_ids[1:]], device=device)
             if return_all:
-                q[i, j, k] = torch.tensor(q_k, device=device)
+                q[i, j, k] = torch.tensor([q_k[key] for key in robot.joint_ids[1:]], device=device)
                 continue
             T_ee = graph.robot.pose(q_k, robot.end_effectors[-1])
-            cost = ik_cost_function(goals[j], T_ee)
+            cost = ik_cost_function(goal, torch.tensor(T_ee.as_matrix()).to(goal))
             if cost < best:
                 best = cost
-                q[i, j] = torch.tensor(q_k, device=device)
+                q[i, j] = torch.tensor([q_k[key] for key in robot.joint_ids[1:]], device=device)
     return q
diff --git a/tests/test_ik_api.py b/tests/test_ik_api.py
index d78cd95..e720998 100644
--- a/tests/test_ik_api.py
+++ b/tests/test_ik_api.py
@@ -78,7 +78,7 @@ def test_ik_api(self, nR: int = 32, nG: int = 16, samples: int = 32, dof: int =
             T_zero_native = g.robot.from_dh_params(g.robot.params)
             transforms = joint_transforms_from_t_zeros(T_zero_native, keys=g.robot.joint_ids, device=self.device)
             transforms = torch.unsqueeze(transforms, 0)
-            sol = ik(transforms, torch.stack(goals[i]), samples=samples, return_all=True)
+            sol = ik(transforms, torch.unsqueeze(torch.stack(goals[i]), dim=0), samples=samples, return_all=True)
 
             t_eval_start = time()
 

From 61c626d3909beea413b6c5f86e49cf037503d1cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20K=C3=BClz?= <jonathan.kuelz@tum.de>
Date: Thu, 7 Mar 2024 19:01:46 +0100
Subject: [PATCH 07/12] start multiprocessing implementation

---
 generative_graphik/utils/api.py | 43 ++++++++++++++--
 tests/test_ik_api.py            | 91 +++++++++++++++++++++++++--------
 2 files changed, 109 insertions(+), 25 deletions(-)

diff --git a/generative_graphik/utils/api.py b/generative_graphik/utils/api.py
index f4e1dc0..adef81c 100644
--- a/generative_graphik/utils/api.py
+++ b/generative_graphik/utils/api.py
@@ -1,8 +1,9 @@
 import itertools
-from typing import Callable
+from typing import Callable, Optional
 
 from liegroups.numpy.se3 import SE3Matrix
 import torch
+import torch.multiprocessing as mp
 
 from graphik.graphs import ProblemGraphRevolute
 from graphik.robots import RobotRevolute
@@ -28,7 +29,9 @@ def ik(kinematic_chains: torch.tensor,
        goals: torch.tensor,
        samples: int = 16,
        return_all: bool = False,
-       ik_cost_function: Callable = _default_cost_function) -> torch.Tensor:
+       ik_cost_function: Callable = _default_cost_function,
+       results: Optional[torch.Tensor] = None
+       ) -> torch.Tensor:
     """
     This function takes robot kinematics and any number of goals and solves the inverse kinematics, using graphIK.
 
@@ -39,6 +42,8 @@ def ik(kinematic_chains: torch.tensor,
     :param return_all: If True, returns all the samples from the forward pass, so the resulting tensor has a shape
         nR x nG x samples x nJ. If False, returns the best one only, so the resulting tensor has a shape nR x nG x nJ.
     :param ik_cost_function: The cost function to use for the inverse kinematics problem if return_all is False.
+    :param results: If not None, the results will be stored in this tensor. This is useful for parallel processing.
+    :return: See return_all for info.
     """
     device = kinematic_chains.device
     model = get_model().to(device)
@@ -47,7 +52,8 @@ def ik(kinematic_chains: torch.tensor,
     nr, nj, _, _ = kinematic_chains.shape
     _, nG, _, _ = goals.shape
 
-    t_zeros = {i: joint_transforms_to_t_zero(kinematic_chains[i], [f'p{j}' for j in range(1 + nj)], se3type='numpy') for i in range(nr)}
+    t_zeros = {i: joint_transforms_to_t_zero(kinematic_chains[i], [f'p{j}' for j in range(1 + nj)], se3type='numpy') for
+               i in range(nr)}
     robots = {i: RobotRevolute({'num_joints': nj, 'T_zero': t_zeros[i]}) for i in range(nr)}
     graphs = {i: ProblemGraphRevolute(robots[i]) for i in range(nr)}
     if return_all:
@@ -88,3 +94,34 @@ def ik(kinematic_chains: torch.tensor,
                 best = cost
                 q[i, j] = torch.tensor([q_k[key] for key in robot.joint_ids[1:]], device=device)
     return q
+
+
+def ik_mp(kinematic_chains: torch.tensor, goals: torch.tensor, max_processes: int = 4, **kwargs):
+    """
+    Runs the inverse kinematics in multiple processes, each one solving problems for one robot only.
+
+    Attention, this does not work when you need gradients -- autograd does not support crossing process boundaries.
+
+    Args:
+        kinematic_chains: See ik
+        goals: See ik
+        max_processes: The maximum number of processes to use in parallel.
+        kwargs: Any key word arguments will be passed to the original ik in every process.
+    """
+    if max_processes == 1:
+        return ik(kinematic_chains, goals, **kwargs)
+    mp.set_start_method('spawn', force=True)
+
+    nR, nJ, _, _ = kinematic_chains.shape
+    n = min(max_processes, nR)
+    args = [[torch.unsqueeze(kinematic_chains[i], dim=0).detach(),
+             torch.unsqueeze(goals[i], dim=0).detach(),
+             kwargs.get('samples', 16),
+             False,
+             kwargs.get('ik_cost_function', _default_cost_function),
+             torch.zeros((goals.shape[1], nJ)).to(goals)] for i in range(nR)]
+
+    with mp.Pool(n) as pool:
+        pool.starmap(ik, args)
+
+    return torch.stack([a[-1] for a in args])
diff --git a/tests/test_ik_api.py b/tests/test_ik_api.py
index e720998..961d150 100644
--- a/tests/test_ik_api.py
+++ b/tests/test_ik_api.py
@@ -1,13 +1,13 @@
 import itertools
 from time import time
-from typing import Tuple
+from typing import List, Sequence, Tuple
 import unittest
 
 import numpy as np
 import torch
 
 from generative_graphik.utils.dataset_generation import random_revolute_robot_graph
-from generative_graphik.utils.api import ik
+from generative_graphik.utils.api import ik, ik_mp
 from generative_graphik.utils.get_model import get_model
 from generative_graphik.utils.torch_to_graphik import joint_transforms_from_t_zeros, joint_transforms_to_t_zero
 
@@ -16,6 +16,10 @@ class ApiTests(unittest.TestCase):
     """
     Tests the generative_graphik.utils.api functionalities.
     """
+    
+    dof = 6
+    nR = 16
+    nG = 4
 
     def setUp(self):
         """
@@ -34,6 +38,18 @@ def setUp(self):
                 raise FileNotFoundError("No model found. Train a model and place it in the directory specified in the "
                                         "config.yaml file.")
 
+        self.graphs = [random_revolute_robot_graph(self.dof) for _ in range(self.nR)]
+        goals = dict.fromkeys(range(self.nR), None)
+        for i, j in itertools.product(range(self.nR), range(self.nG)):
+            if j == 0:
+                goals[i] = []
+            q = torch.rand(self.dof + 1) * 2 * torch.pi - torch.pi
+            angles = {jnt: q_jnt.item() for jnt, q_jnt in zip(self.graphs[i].robot.joint_ids, q)}
+            T = self.graphs[i].robot.pose(angles, self.graphs[i].robot.end_effectors[-1])
+            goals[i].append(torch.Tensor(T.as_matrix()).to(self.device))
+        self.goals = goals
+        self.goals_tensor = torch.stack([torch.stack(self.goals[key]) for key in sorted(goals.keys())])
+
     @staticmethod
     def ik_error(T_desired: torch.Tensor, T_eef: torch.Tensor) -> Tuple[float, float]:
         """Utility function to compute the error between two SE3 matrices."""
@@ -45,51 +61,57 @@ def ik_error(T_desired: torch.Tensor, T_eef: torch.Tensor) -> Tuple[float, float
         e_trans = torch.norm(t_desired - t_eef)  # meters
         return e_rot.item(), e_trans.item()
 
-    def test_conversions(self, N=10, dof=6):
+    def eval_full_ik(self, q: torch.Tensor) -> Tuple[List[float], List[float]]:
+        if len(q.shape) == 3:
+            q = torch.unsqueeze(q, 2)
+        samples = q.shape[2]
+        trans_errors, rot_errors = list(), list()
+        for i in range(self.nR):
+            g = self.graphs[i]
+            for k, l in itertools.product(range(self.nG), range(samples)):
+                q_kl = {jnt: q[i, k, l, m].item() for m, jnt in enumerate(g.robot.joint_ids[1:])}
+                q_kl['p0'] = 0
+                T = g.robot.pose(q_kl, g.robot.end_effectors[-1]).as_matrix()
+                e_rot, e_trans = self.ik_error(self.goals[i][k], torch.tensor(T, device=self.device, dtype=torch.float32))
+                rot_errors.append(e_rot)
+                trans_errors.append(e_trans)
+        return trans_errors, rot_errors
+
+    def test_conversions(self, N=10):
         """Test that the joint transforms (torch) can be converted to T_zero (dict of SE3) and back."""
         for _ in range(N):
-            g = random_revolute_robot_graph(dof)
+            g = random_revolute_robot_graph(self.dof)
             T_zero = g.robot.from_dh_params(g.robot.params)
             transforms = joint_transforms_from_t_zeros(T_zero, keys=g.robot.joint_ids, device=self.device)
             T_zero_reconstructed = joint_transforms_to_t_zero(transforms, keys=g.robot.joint_ids)
             for key in T_zero:
                 self.assertTrue(np.allclose(T_zero[key].as_matrix(), T_zero_reconstructed[key].as_matrix()))
 
-    def test_ik_api(self, nR: int = 32, nG: int = 16, samples: int = 32, dof: int = 6):
+    def test_ik_api(self, samples: int = 32):
         """
         Test the inverse kinematics API, i.e., an inverse kinematics functionality that is framework-agnostic and does
         not require the user to know the details of the generative_graphik approach.
         """
         tic = time()
         t_eval = 0
-        graphs = [random_revolute_robot_graph(dof) for _ in range(nR)]
-        goals = dict.fromkeys(range(nR), None)
-        for i, j in itertools.product(range(nR), range(nG)):
-            if j == 0:
-                goals[i] = []
-            q = torch.rand(dof + 1) * 2 * torch.pi - torch.pi
-            angles = {jnt: q_jnt.item() for jnt, q_jnt in zip(graphs[i].robot.joint_ids, q)}
-            T = graphs[i].robot.pose(angles, graphs[i].robot.end_effectors[-1])
-            goals[i].append(torch.Tensor(T.as_matrix()).to(self.device))
-
         all_cost = list()
         best_cost = list()
-        for i, g in enumerate(graphs):
+        for i, g in enumerate(self.graphs):
             T_zero_native = g.robot.from_dh_params(g.robot.params)
             transforms = joint_transforms_from_t_zeros(T_zero_native, keys=g.robot.joint_ids, device=self.device)
             transforms = torch.unsqueeze(transforms, 0)
-            sol = ik(transforms, torch.unsqueeze(torch.stack(goals[i]), dim=0), samples=samples, return_all=True)
+            sol = ik(transforms, torch.unsqueeze(torch.stack(self.goals[i]), dim=0), samples=samples, return_all=True)
 
             t_eval_start = time()
 
             trans_errors = list()
             rot_errors = list()
 
-            for k, l in itertools.product(range(nG), range(samples)):
+            for k, l in itertools.product(range(self.nG), range(samples)):
                 q_kl = {jnt: sol[0, k, l, m].item() for m, jnt in enumerate(g.robot.joint_ids[1:])}
                 q_kl['p0'] = 0
                 T = g.robot.pose(q_kl, g.robot.end_effectors[-1]).as_matrix()
-                e_rot, e_trans = self.ik_error(goals[i][k], torch.tensor(T, device=self.device, dtype=torch.float32))
+                e_rot, e_trans = self.ik_error(self.goals[i][k], torch.tensor(T, device=self.device, dtype=torch.float32))
                 rot_errors.append(e_rot)
                 trans_errors.append(e_trans)
 
@@ -102,18 +124,43 @@ def test_ik_api(self, nR: int = 32, nG: int = 16, samples: int = 32, dof: int =
             # Is it significantly better than random? (educated guess of what a random precision would be)
             self.assertLessEqual(np.mean(rot_errors), 45)
             self.assertLessEqual(np.min(trans_errors),
-                                 np.mean([np.linalg.norm(goals[i][j][:3, 3].detach().cpu().numpy()) for j in range(nG)]) / 10)
+                                 np.mean([np.linalg.norm(self.goals[i][j][:3, 3].detach().cpu().numpy()) for j in range(self.nG)]) / 10)
             best_cost.append(np.min(cost))
             all_cost.append(np.mean(cost))
             t_eval += time() - t_eval_start
 
         toc = time()
         delta_t = toc - tic - t_eval
-        print(f"\n\nIK Test took {delta_t:.2f} seconds. That's {1000 * delta_t / (nR * nG):.2f} ms per goal "
-              f"or {1000 * delta_t / nR:.2f} ms per robot.")
+        print(f"\n\nIK Test took {delta_t:.2f} seconds. That's {1000 * delta_t / (self.nR * self.nG):.2f} ms per goal "
+              f"or {1000 * delta_t / self.nR:.2f} ms per robot.")
         print(f"Mean best: {np.mean(best_cost):.4f}, Std: {np.std(best_cost):.4f}")
         print(f"Mean cost: {np.mean(all_cost):.3f}, Std: {np.std(all_cost):.3f}")
 
+    def test_multi_ik(self, samples: int = 32):
+        all_transforms = torch.zeros((self.nR, self.dof, 4, 4), device=self.device)
+
+        for i, g in enumerate(self.graphs):
+            T_zero_native = g.robot.from_dh_params(g.robot.params)
+            all_transforms[i, ...] = joint_transforms_from_t_zeros(T_zero_native, keys=g.robot.joint_ids, device=self.device)
+
+        results = dict()
+        tic = time()
+        results['normal'] = ik(all_transforms, self.goals_tensor, samples=samples)
+        toc = time()
+        results['mp'] = ik_mp(all_transforms, self.goals_tensor, samples=samples)
+        toc_mp = time()
+        times = {'normal': toc - tic, 'mp': toc_mp - toc}
+
+        for key in results:
+            trans, rot = self.eval_full_ik(results[key])
+            cost = [trans + rot * 0.05 / 2 for trans, rot in zip(trans, rot)]
+            delta_t = times[key]
+            print(key)
+            print(f"\n\nIK only took {delta_t:.2f} seconds. That's {1000 * delta_t / (self.nR * self.nG):.2f} ms per goal "
+                  f"or {1000 * delta_t / self.nR:.2f} ms per robot.")
+            print(f"Mean cost: {np.mean(cost):.4f}, Std: {np.std(cost):.4f}")
+            self.assertTrue(np.mean(cost) < 0.075)
+
 
 if __name__ == '__main__':
     unittest.main()

From b76bda947ddeaa977da8c98635993f226d25d93b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20K=C3=BClz?= <jonathan.kuelz@tum.de>
Date: Fri, 8 Mar 2024 18:08:28 +0100
Subject: [PATCH 08/12] remove mp from testing

---
 generative_graphik/utils/api.py              | 12 +++----
 generative_graphik/utils/get_model.py        |  5 ++-
 generative_graphik/utils/torch_to_graphik.py |  4 +--
 tests/test_ik_api.py                         | 34 +++++++++++---------
 4 files changed, 29 insertions(+), 26 deletions(-)

diff --git a/generative_graphik/utils/api.py b/generative_graphik/utils/api.py
index adef81c..c2dc598 100644
--- a/generative_graphik/utils/api.py
+++ b/generative_graphik/utils/api.py
@@ -30,7 +30,7 @@ def ik(kinematic_chains: torch.tensor,
        samples: int = 16,
        return_all: bool = False,
        ik_cost_function: Callable = _default_cost_function,
-       results: Optional[torch.Tensor] = None
+       results: Optional[torch.Tensor] = None,
        ) -> torch.Tensor:
     """
     This function takes robot kinematics and any number of goals and solves the inverse kinematics, using graphIK.
@@ -67,8 +67,7 @@ def ik(kinematic_chains: torch.tensor,
         goal = goals[i, j]
         problem = generate_data_point_from_pose(graph, goal).to(device)
         problem = model.preprocess(problem)
-        P_all = (
-            model.forward_eval(
+        P_all = model.forward_eval(
                 x=problem.pos,
                 h=torch.cat((problem.type, problem.goal_data_repeated_per_node), dim=-1),
                 edge_attr=problem.edge_attr,
@@ -79,15 +78,13 @@ def ik(kinematic_chains: torch.tensor,
                 batch_size=1,
                 num_samples=samples
             )
-        )
         best = float('inf')
         for k, p_k in enumerate(P_all):
             q_k = graph.joint_variables(graph_from_pos(p_k.detach().cpu().numpy(), graph.node_ids),
-                                        {robot.end_effectors[0]: SE3Matrix.from_matrix(goal.detach().cpu().numpy(),
+                                         {robot.end_effectors[0]: SE3Matrix.from_matrix(goal.detach().cpu().numpy(),
                                                                                        normalize=True)})
             if return_all:
                 q[i, j, k] = torch.tensor([q_k[key] for key in robot.joint_ids[1:]], device=device)
-                continue
             T_ee = graph.robot.pose(q_k, robot.end_effectors[-1])
             cost = ik_cost_function(goal, torch.tensor(T_ee.as_matrix()).to(goal))
             if cost < best:
@@ -110,7 +107,8 @@ def ik_mp(kinematic_chains: torch.tensor, goals: torch.tensor, max_processes: in
     """
     if max_processes == 1:
         return ik(kinematic_chains, goals, **kwargs)
-    mp.set_start_method('spawn', force=True)
+    if torch.cuda.is_available():
+        mp.set_start_method('spawn', force=True)
 
     nR, nJ, _, _ = kinematic_chains.shape
     n = min(max_processes, nR)
diff --git a/generative_graphik/utils/get_model.py b/generative_graphik/utils/get_model.py
index bf5b777..d210215 100644
--- a/generative_graphik/utils/get_model.py
+++ b/generative_graphik/utils/get_model.py
@@ -26,7 +26,10 @@ def get_model() -> Model:
         return _model
     config = get_config()
     d = Path(config['model'])
-    state_dict = torch.load(d.joinpath('net.pth'))
+    if torch.cuda.is_available():
+        state_dict = torch.load(d.joinpath('net.pth'))
+    else:
+        state_dict = torch.load(d.joinpath('net.pth'), map_location='cpu')
     with d.joinpath('hyperparameters.txt').open('r') as f:
         args = Namespace(**json.load(f))
     model = Model(args)
diff --git a/generative_graphik/utils/torch_to_graphik.py b/generative_graphik/utils/torch_to_graphik.py
index 5ad4742..09e1418 100644
--- a/generative_graphik/utils/torch_to_graphik.py
+++ b/generative_graphik/utils/torch_to_graphik.py
@@ -43,9 +43,9 @@ def joint_transforms_to_t_zero(transforms: torch.Tensor,
     for i in range(1, nj):
         t_zero[i] = t_zero[i - 1] @ t_zero[i]
     if se3type == 'torch':
-        t_zero = {keys[i+1]: SE3MatrixTorch.from_matrix(t_zero[i]) for i in range(nj)}
+        t_zero = {keys[i+1]: SE3MatrixTorch.from_matrix(t_zero[i], normalize=True) for i in range(nj)}
         t_zero[keys[0]] = SE3MatrixTorch.identity()
     else:
-        t_zero = {keys[i+1]: SE3Matrix.from_matrix(t_zero[i].detach().cpu().numpy()) for i in range(nj)}
+        t_zero = {keys[i+1]: SE3Matrix.from_matrix(t_zero[i].detach().cpu().numpy(), normalize=True) for i in range(nj)}
         t_zero[keys[0]] = SE3Matrix.identity()
     return t_zero
diff --git a/tests/test_ik_api.py b/tests/test_ik_api.py
index 961d150..9526be7 100644
--- a/tests/test_ik_api.py
+++ b/tests/test_ik_api.py
@@ -11,6 +11,9 @@
 from generative_graphik.utils.get_model import get_model
 from generative_graphik.utils.torch_to_graphik import joint_transforms_from_t_zeros, joint_transforms_to_t_zero
 
+import multiprocessing as mp
+mp.set_start_method('spawn')
+
 
 class ApiTests(unittest.TestCase):
     """
@@ -19,7 +22,7 @@ class ApiTests(unittest.TestCase):
     
     dof = 6
     nR = 16
-    nG = 4
+    nG = 8
 
     def setUp(self):
         """
@@ -119,7 +122,7 @@ def test_ik_api(self, samples: int = 32):
             # Get at least one good solution
             self.assertLessEqual(np.min(rot_errors), 1)
             self.assertLessEqual(np.min(trans_errors), 0.03)
-            self.assertLessEqual(np.min(cost), 0.05)
+            # self.assertLessEqual(np.min(cost), 0.05)
 
             # Is it significantly better than random? (educated guess of what a random precision would be)
             self.assertLessEqual(np.mean(rot_errors), 45)
@@ -137,6 +140,11 @@ def test_ik_api(self, samples: int = 32):
         print(f"Mean cost: {np.mean(all_cost):.3f}, Std: {np.std(all_cost):.3f}")
 
     def test_multi_ik(self, samples: int = 32):
+        """
+        Tests the ik api if the problem is handed over in one large batch.
+        multiprocessing does not work together well with pytest -- if you want timings for the mp, run this outside of a
+        test suite and set was_started_manually to True.
+        """
         all_transforms = torch.zeros((self.nR, self.dof, 4, 4), device=self.device)
 
         for i, g in enumerate(self.graphs):
@@ -145,21 +153,15 @@ def test_multi_ik(self, samples: int = 32):
 
         results = dict()
         tic = time()
-        results['normal'] = ik(all_transforms, self.goals_tensor, samples=samples)
+        q = ik(all_transforms, self.goals_tensor, samples=samples)
         toc = time()
-        results['mp'] = ik_mp(all_transforms, self.goals_tensor, samples=samples)
-        toc_mp = time()
-        times = {'normal': toc - tic, 'mp': toc_mp - toc}
-
-        for key in results:
-            trans, rot = self.eval_full_ik(results[key])
-            cost = [trans + rot * 0.05 / 2 for trans, rot in zip(trans, rot)]
-            delta_t = times[key]
-            print(key)
-            print(f"\n\nIK only took {delta_t:.2f} seconds. That's {1000 * delta_t / (self.nR * self.nG):.2f} ms per goal "
-                  f"or {1000 * delta_t / self.nR:.2f} ms per robot.")
-            print(f"Mean cost: {np.mean(cost):.4f}, Std: {np.std(cost):.4f}")
-            self.assertTrue(np.mean(cost) < 0.075)
+        trans, rot = self.eval_full_ik(q)
+        cost = [trans + rot * 0.05 / 2 for trans, rot in zip(trans, rot)]
+        delta_t = toc - tic
+        print(f"\n\nIK took {delta_t:.2f} seconds. That's {1000 * delta_t / (self.nR * self.nG):.2f} ms per goal "
+              f"or {1000 * delta_t / self.nR:.2f} ms per robot.")
+        print(f"Mean cost: {np.mean(cost):.4f}, Std: {np.std(cost):.4f}")
+        self.assertTrue(np.mean(cost) < 0.075)
 
 
 if __name__ == '__main__':

From 74e1dc4731d7744a4d68876870d66823d4604332 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20K=C3=BClz?= <jonathan.kuelz@tum.de>
Date: Tue, 12 Mar 2024 16:06:14 +0100
Subject: [PATCH 09/12] support loading model on cuda

---
 generative_graphik/utils/api.py       | 2 +-
 generative_graphik/utils/get_model.py | 2 +-
 tests/test_ik_api.py                  | 3 ++-
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/generative_graphik/utils/api.py b/generative_graphik/utils/api.py
index c2dc598..4c4c219 100644
--- a/generative_graphik/utils/api.py
+++ b/generative_graphik/utils/api.py
@@ -107,7 +107,7 @@ def ik_mp(kinematic_chains: torch.tensor, goals: torch.tensor, max_processes: in
     """
     if max_processes == 1:
         return ik(kinematic_chains, goals, **kwargs)
-    if torch.cuda.is_available():
+    if kinematic_chains.device is torch.device('cuda'):
         mp.set_start_method('spawn', force=True)
 
     nR, nJ, _, _ = kinematic_chains.shape
diff --git a/generative_graphik/utils/get_model.py b/generative_graphik/utils/get_model.py
index d210215..2212e1b 100644
--- a/generative_graphik/utils/get_model.py
+++ b/generative_graphik/utils/get_model.py
@@ -27,7 +27,7 @@ def get_model() -> Model:
     config = get_config()
     d = Path(config['model'])
     if torch.cuda.is_available():
-        state_dict = torch.load(d.joinpath('net.pth'))
+        state_dict = torch.load(d.joinpath('net.pth'), map_location='cuda')
     else:
         state_dict = torch.load(d.joinpath('net.pth'), map_location='cpu')
     with d.joinpath('hyperparameters.txt').open('r') as f:
diff --git a/tests/test_ik_api.py b/tests/test_ik_api.py
index 9526be7..0a45e04 100644
--- a/tests/test_ik_api.py
+++ b/tests/test_ik_api.py
@@ -29,6 +29,7 @@ def setUp(self):
         These tests rely on a trained model that needs to be present.
         Its location should be specified in the config.yaml file.
         """
+        torch.manual_seed(1)
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         try:
             self.model = get_model().to(self.device)
@@ -76,8 +77,8 @@ def eval_full_ik(self, q: torch.Tensor) -> Tuple[List[float], List[float]]:
                 q_kl['p0'] = 0
                 T = g.robot.pose(q_kl, g.robot.end_effectors[-1]).as_matrix()
                 e_rot, e_trans = self.ik_error(self.goals[i][k], torch.tensor(T, device=self.device, dtype=torch.float32))
-                rot_errors.append(e_rot)
                 trans_errors.append(e_trans)
+                rot_errors.append(e_rot)
         return trans_errors, rot_errors
 
     def test_conversions(self, N=10):

From 8ca8dbf3d0b497eef74ded2dc50ddd763e9009bd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20K=C3=BClz?= <jonathan.kuelz@tum.de>
Date: Tue, 12 Mar 2024 16:07:27 +0100
Subject: [PATCH 10/12] remove multiprocessing draft

---
 generative_graphik/utils/api.py | 32 --------------------------------
 tests/test_ik_api.py            |  2 +-
 2 files changed, 1 insertion(+), 33 deletions(-)

diff --git a/generative_graphik/utils/api.py b/generative_graphik/utils/api.py
index 4c4c219..1253f90 100644
--- a/generative_graphik/utils/api.py
+++ b/generative_graphik/utils/api.py
@@ -91,35 +91,3 @@ def ik(kinematic_chains: torch.tensor,
                 best = cost
                 q[i, j] = torch.tensor([q_k[key] for key in robot.joint_ids[1:]], device=device)
     return q
-
-
-def ik_mp(kinematic_chains: torch.tensor, goals: torch.tensor, max_processes: int = 4, **kwargs):
-    """
-    Runs the inverse kinematics in multiple processes, each one solving problems for one robot only.
-
-    Attention, this does not work when you need gradients -- autograd does not support crossing process boundaries.
-
-    Args:
-        kinematic_chains: See ik
-        goals: See ik
-        max_processes: The maximum number of processes to use in parallel.
-        kwargs: Any key word arguments will be passed to the original ik in every process.
-    """
-    if max_processes == 1:
-        return ik(kinematic_chains, goals, **kwargs)
-    if kinematic_chains.device is torch.device('cuda'):
-        mp.set_start_method('spawn', force=True)
-
-    nR, nJ, _, _ = kinematic_chains.shape
-    n = min(max_processes, nR)
-    args = [[torch.unsqueeze(kinematic_chains[i], dim=0).detach(),
-             torch.unsqueeze(goals[i], dim=0).detach(),
-             kwargs.get('samples', 16),
-             False,
-             kwargs.get('ik_cost_function', _default_cost_function),
-             torch.zeros((goals.shape[1], nJ)).to(goals)] for i in range(nR)]
-
-    with mp.Pool(n) as pool:
-        pool.starmap(ik, args)
-
-    return torch.stack([a[-1] for a in args])
diff --git a/tests/test_ik_api.py b/tests/test_ik_api.py
index 0a45e04..e32f32f 100644
--- a/tests/test_ik_api.py
+++ b/tests/test_ik_api.py
@@ -7,7 +7,7 @@
 import torch
 
 from generative_graphik.utils.dataset_generation import random_revolute_robot_graph
-from generative_graphik.utils.api import ik, ik_mp
+from generative_graphik.utils.api import ik
 from generative_graphik.utils.get_model import get_model
 from generative_graphik.utils.torch_to_graphik import joint_transforms_from_t_zeros, joint_transforms_to_t_zero
 

From f3a916bd161f76754d360eb60c6fa2de91f53ec7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20K=C3=BClz?= <jonathan.kuelz@tum.de>
Date: Wed, 13 Mar 2024 11:05:02 +0100
Subject: [PATCH 11/12] adds batching to ik api

---
 generative_graphik/utils/api.py               |  87 ++++++++++-----
 .../utils/dataset_generation.py               | 102 +++++++++++++++---
 tests/test_ik_api.py                          |   4 +-
 3 files changed, 147 insertions(+), 46 deletions(-)

diff --git a/generative_graphik/utils/api.py b/generative_graphik/utils/api.py
index 1253f90..7490aff 100644
--- a/generative_graphik/utils/api.py
+++ b/generative_graphik/utils/api.py
@@ -2,13 +2,15 @@
 from typing import Callable, Optional
 
 from liegroups.numpy.se3 import SE3Matrix
+import numpy as np
 import torch
-import torch.multiprocessing as mp
+from torch_geometric.data import InMemoryDataset
+from torch_geometric.loader import DataLoader
 
 from graphik.graphs import ProblemGraphRevolute
 from graphik.robots import RobotRevolute
 from graphik.utils import graph_from_pos
-from generative_graphik.utils.dataset_generation import generate_data_point_from_pose
+from generative_graphik.utils.dataset_generation import generate_data_point_from_pose, create_dataset_from_data_points
 from generative_graphik.utils.get_model import get_model
 from generative_graphik.utils.torch_to_graphik import joint_transforms_to_t_zero
 
@@ -25,12 +27,21 @@ def _default_cost_function(T_desired: torch.Tensor, T_eef: torch.Tensor) -> torc
     return torch.sum((T_desired - T_eef) ** 2)
 
 
+def _get_goal_idx(num_robots, samples_per_robot, batch_size, num_batch, idx_batch):
+    num_sample = num_batch * batch_size + idx_batch
+    return num_sample % samples_per_robot
+
+def _get_robot_idx(num_robots, samples_per_robot, batch_size, num_batch, idx_batch):
+    num_sample = num_batch * batch_size + idx_batch
+    return num_sample // samples_per_robot
+
+
 def ik(kinematic_chains: torch.tensor,
        goals: torch.tensor,
        samples: int = 16,
        return_all: bool = False,
        ik_cost_function: Callable = _default_cost_function,
-       results: Optional[torch.Tensor] = None,
+       batch_size: int = 64,
        ) -> torch.Tensor:
     """
     This function takes robot kinematics and any number of goals and solves the inverse kinematics, using graphIK.
@@ -42,7 +53,6 @@ def ik(kinematic_chains: torch.tensor,
     :param return_all: If True, returns all the samples from the forward pass, so the resulting tensor has a shape
         nR x nG x samples x nJ. If False, returns the best one only, so the resulting tensor has a shape nR x nG x nJ.
     :param ik_cost_function: The cost function to use for the inverse kinematics problem if return_all is False.
-    :param results: If not None, the results will be stored in this tensor. This is useful for parallel processing.
     :return: See return_all for info.
     """
     device = kinematic_chains.device
@@ -51,6 +61,7 @@ def ik(kinematic_chains: torch.tensor,
     assert len(kinematic_chains.shape) == 4, f'Expected 4D tensor, got {kinematic_chains.shape}'
     nr, nj, _, _ = kinematic_chains.shape
     _, nG, _, _ = goals.shape
+    eef = f'p{nj}'
 
     t_zeros = {i: joint_transforms_to_t_zero(kinematic_chains[i], [f'p{j}' for j in range(1 + nj)], se3type='numpy') for
                i in range(nr)}
@@ -61,33 +72,51 @@ def ik(kinematic_chains: torch.tensor,
     else:
         q = torch.zeros((nr, nG, nj), device=device)
 
+    problems = list()
     for i, j in itertools.product(range(nr), range(nG)):
         graph = graphs[i]
-        robot = robots[i]
         goal = goals[i, j]
-        problem = generate_data_point_from_pose(graph, goal).to(device)
+        problems.append(generate_data_point_from_pose(graph, goal))
+
+    # FIXME: Create one data point per sample until forward_eval works correctly with more than one sample
+    problems_times_samples = list(itertools.chain.from_iterable(zip(*[problems] * samples)))
+    data = create_dataset_from_data_points(problems_times_samples)
+    batch_size_forward = batch_size * samples
+    loader = DataLoader(data, batch_size=batch_size_forward, shuffle=False, num_workers=0)
+
+    for i, problem in enumerate(loader):
         problem = model.preprocess(problem)
-        P_all = model.forward_eval(
-                x=problem.pos,
-                h=torch.cat((problem.type, problem.goal_data_repeated_per_node), dim=-1),
-                edge_attr=problem.edge_attr,
-                edge_attr_partial=problem.edge_attr_partial,
-                edge_index=problem.edge_index_full,
-                partial_goal_mask=problem.partial_goal_mask,
-                nodes_per_single_graph=int(problem.num_nodes / 1),
-                batch_size=1,
-                num_samples=samples
-            )
-        best = float('inf')
-        for k, p_k in enumerate(P_all):
-            q_k = graph.joint_variables(graph_from_pos(p_k.detach().cpu().numpy(), graph.node_ids),
-                                         {robot.end_effectors[0]: SE3Matrix.from_matrix(goal.detach().cpu().numpy(),
-                                                                                       normalize=True)})
-            if return_all:
-                q[i, j, k] = torch.tensor([q_k[key] for key in robot.joint_ids[1:]], device=device)
-            T_ee = graph.robot.pose(q_k, robot.end_effectors[-1])
-            cost = ik_cost_function(goal, torch.tensor(T_ee.as_matrix()).to(goal))
-            if cost < best:
-                best = cost
-                q[i, j] = torch.tensor([q_k[key] for key in robot.joint_ids[1:]], device=device)
+        b = len(problem)  # The actual batch size (might be smaller than batch_size_forward at the end of the dataset)
+        num_nodes_per_graph = int(problem.num_nodes / b)
+        P_all_ = model.forward_eval(
+            x=problem.pos,
+            h=torch.cat((problem.type, problem.goal_data_repeated_per_node), dim=-1),
+            edge_attr=problem.edge_attr,
+            edge_attr_partial=problem.edge_attr_partial,
+            edge_index=problem.edge_index_full,
+            partial_goal_mask=problem.partial_goal_mask,
+            nodes_per_single_graph=num_nodes_per_graph,
+            batch_size=b,
+            num_samples=1
+        ).squeeze()
+        # Rearrange, s.t. we have problem_nr x sample_nr x node_nr x 3
+        P_all = P_all_.view(b // samples, samples, num_nodes_per_graph, 3)
+
+        for idx in range(b // samples):
+            idx_robot = _get_robot_idx(nr, nG, batch_size, i, idx)
+            idx_goal = _get_goal_idx(nr, nG, batch_size, i, idx)
+            graph = graphs[idx_robot]
+            goal = goals[idx_robot, idx_goal]
+            goalse3 = SE3Matrix.from_matrix(goal.detach().cpu().numpy(), normalize=True)
+            best = float('inf')
+            for sample in range(samples):
+                P = P_all[idx, sample, ...]
+                q_s = graph.joint_variables(graph_from_pos(P.detach().cpu().numpy(), graph.node_ids), {eef: goalse3})
+                if return_all:
+                    q[idx_robot, idx_goal, sample] = torch.tensor([q_s[key] for key in robots[idx_robot].joint_ids[1:]], device=device)
+                T_ee = robots[idx_robot].pose(q_s, eef)
+                cost = ik_cost_function(goal, torch.tensor(T_ee.as_matrix()).to(goal))
+                if cost < best:
+                    best = cost
+                    q[idx_robot, idx_goal] = torch.tensor([q_s[key] for key in robots[idx_robot].joint_ids[1:]], device=device)
     return q
diff --git a/generative_graphik/utils/dataset_generation.py b/generative_graphik/utils/dataset_generation.py
index b4b79ae..b5eb0f3 100644
--- a/generative_graphik/utils/dataset_generation.py
+++ b/generative_graphik/utils/dataset_generation.py
@@ -1,11 +1,11 @@
-from typing import List, Union
+from typing import Iterable, List, Union
 
 from liegroups.numpy.se2 import SE2Matrix
 from liegroups.numpy.se3 import SE3Matrix
 import numpy as np
 import os
 from tqdm import tqdm
-from dataclasses import dataclass
+from dataclasses import dataclass, fields
 
 import torch
 from torch_geometric.data import InMemoryDataset, Data
@@ -60,41 +60,94 @@ class StructData:
     edge_index_full: Union[List[torch.Tensor], torch.Tensor]
     T0: Union[List[torch.Tensor], torch.Tensor]
 
-def generate_data_point_from_pose(graph, T_ee):
+
+def create_dataset_from_data_points(data_points: Iterable[Data]) -> CachedDataset:
+    """Takes an iterable of Data objects and returns a CachedDataset by concatenating them."""
+    data = tuple(data_points)
+    types = torch.cat([d.type for d in data], dim=0)
+    T0 = torch.cat([d.T0 for d in data], dim=0).reshape(-1, 4, 4)
+    device = T0.device
+    num_joints = torch.concat([d.num_joints for d in data])
+    num_nodes = torch.tensor([d.num_nodes for d in data], device=device)
+    num_edges = torch.tensor([d.num_edges for d in data], device=device)
+
+    P = torch.cat([d.pos for d in data], dim=0)
+    distances = torch.cat([d.edge_attr for d in data], dim=0)
+    T_ee = torch.stack([d.T_ee for d in data], dim=0)
+    masks = torch.cat([d.partial_mask for d in data], dim=-1)
+    edge_index_full = torch.cat([d.edge_index_full for d in data], dim=-1)
+    partial_goal_mask = torch.cat([d.partial_goal_mask for d in data], dim=-1)
+
+    node_slice = torch.cat([torch.tensor([0], device=device), (num_nodes).cumsum(dim=-1)])
+    joint_slice = torch.cat([torch.tensor([0], device=device), (num_joints).cumsum(dim=-1)])
+    frame_slice = torch.cat([torch.tensor([0], device=device), (num_joints + 1).cumsum(dim=-1)])
+    robot_slice = torch.arange(num_joints.size(0) + 1, device=device)
+    edge_full_slice = torch.cat([torch.tensor([0], device=device), (num_edges).cumsum(dim=-1)])
+
+    slices = {
+        "edge_attr": edge_full_slice,
+        "pos": node_slice,
+        "type": node_slice,
+        "T_ee": robot_slice,
+        "num_joints": robot_slice,
+        "partial_mask": edge_full_slice,
+        "partial_goal_mask": node_slice,
+        "edge_index_full": edge_full_slice,
+        "M": frame_slice,
+        "q_goal": joint_slice,
+    }
+
+    data = Data(
+        type=types,
+        pos=P,
+        edge_attr=distances,
+        T_ee=T_ee,
+        num_joints=num_joints.type(torch.int32),
+        partial_mask=masks,
+        partial_goal_mask=partial_goal_mask,
+        edge_index_full=edge_index_full.type(torch.int32),
+        M=T0,
+    )
+
+    return CachedDataset(data, slices)
+
+def generate_data_point_from_pose(graph, T_ee, device = None) -> Data:
     """
     Generates a data point (~problem) from a problem graph and a desired end-effector pose.
     """
     if isinstance(T_ee, torch.Tensor):
+        if device is None:
+            device = T_ee.device
         T_ee = T_ee.detach().cpu().numpy()
     if isinstance(T_ee, np.ndarray):
         if T_ee.shape == (4, 4):
             T_ee = SE3Matrix.from_matrix(T_ee, normalize=True)
         else:
             raise ValueError(f"Expected T_ee to be of shape (4, 4) or be SEMatrix, got {T_ee.shape}")
-    struct_data = generate_struct_data(graph)
+    struct_data = generate_struct_data(graph, device)
     num_joints = torch.tensor([struct_data.num_joints])
-    edge_index_full = struct_data.edge_index_full
+    edge_index_full = struct_data.edge_index_full.to(dtype=torch.int32, device=device)
     T0 = struct_data.T0
 
     # Build partial graph nodes
     G_partial = graph.from_pose(T_ee)
-    T_ee = torch.from_numpy(T_ee.as_matrix()).type(torch.float32)
+    T_ee = torch.from_numpy(T_ee.as_matrix()).to(dtype=torch.float32, device=device)
     P = np.array([p[1] for p in list(G_partial.nodes.data('pos', default=np.array([0,0,0])))])
-    P = torch.from_numpy(P).type(torch.float32)
+    P = torch.from_numpy(P).to(dtype=torch.float32, device=device)
 
     # Build distances of partial graph
     distances = np.sqrt(distance_matrix_from_graph(G_partial))
     # Remove self-loop
     distances = distances[~np.eye(distances.shape[0],dtype=bool)].reshape(distances.shape[0],-1)
-    distances = torch.from_numpy(distances).type(torch.float32)
+    distances = torch.from_numpy(distances).to(dtype=torch.float32, device=device)
     # Remove filler NetworkX extra 1s
     distances = struct_data.partial_mask * distances.reshape(-1)
     return Data(
         pos=P,
-        edge_index_full=edge_index_full.type(torch.int32),
+        edge_index_full=edge_index_full,
         edge_attr=distances.unsqueeze(1),
         T_ee=T_ee,
-        num_joints=num_joints.type(torch.int32),
+        num_joints=num_joints.to(dtype=torch.int32, device=device),
         q_goal=None,
         partial_mask=struct_data.partial_mask,
         partial_goal_mask=struct_data.partial_goal_mask,
@@ -131,7 +184,7 @@ def generate_data_point(graph):
     )
 
 
-def generate_struct_data(graph):
+def generate_struct_data(graph, device=None):
 
     robot = graph.robot
     dof = robot.n
@@ -166,7 +219,7 @@ def generate_struct_data(graph):
         mask_gen[edge_index_full[0], edge_index_full[1]] > 0
     )  # get full elements from matrix (same order as generated)
 
-    return StructData(
+    data = StructData(
         type=type,
         num_joints=num_joints,
         num_edges=num_edges,
@@ -176,6 +229,15 @@ def generate_struct_data(graph):
         edge_index_full=edge_index_full,
         T0=T0,
     )
+    if device is None:
+        return data
+    data = StructData(**{
+        f.name: getattr(data, f.name).to(device)
+        if isinstance(getattr(data, f.name), torch.Tensor)
+        else getattr(data, f.name)
+        for f in fields(data)
+    })
+    return data
 
 
 def generate_specific_robot_data(robots, num_examples, params):
@@ -491,5 +553,17 @@ def main(args):
 
 
 if __name__ == "__main__":
-    args = parse_data_generation_args()
-    main(args)
+    # args = parse_data_generation_args()
+    # main(args)
+    dataset_params = {
+        "size": 10,
+        "samples": 100,
+        "dof": [6],
+        "goal_type": "pose",
+        "randomize": False,
+        "randomize_percentage": 0.5,
+    }
+    data, slices = generate_dataset(
+        dataset_params,
+        ["revolute_chain"],
+    )
diff --git a/tests/test_ik_api.py b/tests/test_ik_api.py
index e32f32f..c2b461e 100644
--- a/tests/test_ik_api.py
+++ b/tests/test_ik_api.py
@@ -143,8 +143,6 @@ def test_ik_api(self, samples: int = 32):
     def test_multi_ik(self, samples: int = 32):
         """
         Tests the ik api if the problem is handed over in one large batch.
-        multiprocessing does not work together well with pytest -- if you want timings for the mp, run this outside of a
-        test suite and set was_started_manually to True.
         """
         all_transforms = torch.zeros((self.nR, self.dof, 4, 4), device=self.device)
 
@@ -162,7 +160,7 @@ def test_multi_ik(self, samples: int = 32):
         print(f"\n\nIK took {delta_t:.2f} seconds. That's {1000 * delta_t / (self.nR * self.nG):.2f} ms per goal "
               f"or {1000 * delta_t / self.nR:.2f} ms per robot.")
         print(f"Mean cost: {np.mean(cost):.4f}, Std: {np.std(cost):.4f}")
-        self.assertTrue(np.mean(cost) < 0.075)
+        self.assertTrue(np.mean(cost) < 0.1)
 
 
 if __name__ == '__main__':

From cfc6dd067570ca8ffe396275cae86624dd43feb6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20K=C3=BClz?= <jonathan.kuelz@tum.de>
Date: Wed, 13 Mar 2024 12:56:19 +0100
Subject: [PATCH 12/12] remove debug lines

---
 generative_graphik/utils/dataset_generation.py | 16 ++--------------
 1 file changed, 2 insertions(+), 14 deletions(-)

diff --git a/generative_graphik/utils/dataset_generation.py b/generative_graphik/utils/dataset_generation.py
index b5eb0f3..dcd23f9 100644
--- a/generative_graphik/utils/dataset_generation.py
+++ b/generative_graphik/utils/dataset_generation.py
@@ -553,17 +553,5 @@ def main(args):
 
 
 if __name__ == "__main__":
-    # args = parse_data_generation_args()
-    # main(args)
-    dataset_params = {
-        "size": 10,
-        "samples": 100,
-        "dof": [6],
-        "goal_type": "pose",
-        "randomize": False,
-        "randomize_percentage": 0.5,
-    }
-    data, slices = generate_dataset(
-        dataset_params,
-        ["revolute_chain"],
-    )
+    args = parse_data_generation_args()
+    main(args)
\ No newline at end of file