implement robot_arm

src/inverse_kinematics_sbi/components.py

@@ -76,11 +76,16 @@ def forward_kinematics(self, params):
         return action
 
 
-class ConstAction(ABC):
+class ConstComponent(Component):
 
     def __init__(self, action):
         self.action = action
 
-    def forward_kinematics(self):
-        return self.action
+    def forward_kinematics(self, params):
+        self.check_params(params)
+
+        action = np.zeros(params.shape + (3,))
+        action[..., :] = self.action
+
+        return action
 

src/inverse_kinematics_sbi/proba_robot.py

@@ -19,7 +19,7 @@ def sample_joint(self, n_sample):
 
     def sample_likelihood(self, params, n_sample=1, keep_dims=False):
         # if keep_dims=True and n_sample=1 returns shape: (len(params), 1, n_params)
-        return # forward kinematics on zero
+        return # forward
 
     def sample_posterior(self, observations, n_sample=1, keep_dims=False):
         # if keep_dims=True and n_sample=1 returns shape: (len(params), 1, n_params)

src/inverse_kinematics_sbi/robots.py

@@ -2,29 +2,91 @@
 
 import numpy as np
 
+from src.inverse_kinematics_sbi.components import Component, ConstComponent, SimpleRail, SimpleJoint
+from src.inverse_kinematics_sbi.trigonometry import se2_compose, se2_action
+
 
 class RobotArm(ABC):
 
     def __init__(self, components):
         self.components = components
 
+    def check_components(self):
+        new_components = []
+        for i, component in enumerate(self.components):
+            if isinstance(component, RobotArm):
+                component.check_components()
+                new_components = new_components + component.components
+            elif isinstance(component, Component):
+                new_components = new_components + [component]
+            else:
+                raise ValueError(
+                    f"All components of `self.component` must be of type `RobotArm` or `Component` but component {i} is"
+                    f"of type {type(component)}."
+                )
+
     def get_n_params(self):
-        return  # number of non const components
+        return sum(not isinstance(comp, ConstComponent) for comp in self.components)
+
+    def forward_kinematics(self, params, return_intermediates=False):
+        self.check_components()
+        current_action = np.zeros(params.shape[:-1] + (3,))
+        if return_intermediates:
+            intermediate_actions = np.zeros(params.shape[:-1] + (1+self.get_n_params(), 3))
+        current_index = 0
 
-    def forward_kinematics(self, params):
-        return  # composed action
+        for component in self.components:
+            if return_intermediates and not isinstance(component, ConstComponent):
+                intermediate_actions[..., current_index, :] = current_action
+            component_action = component.forward_kinematics(params[..., current_index])
+            current_action = se2_compose(current_action, component_action)
+            if not isinstance(component, ConstComponent):
+                current_index += 1
 
-    def forward_kinematics_reduced(self, params, start_indices, end_indices):
-        return  # composed reduced action
+        if return_intermediates:
+            intermediate_actions[..., current_index, :] = current_action
+            return intermediate_actions
+        else:
+            return current_action
 
-    def forward(self, params, start_positions=None):
-        return # end_positions
+    def forward(self, params, start_positions=None, return_intermediates=False):
+        if start_positions is None:
+            start_positions = np.zeros(params.shape[:-1] + (2,))
+        if return_intermediates:
+            return se2_action(start_positions[..., None, :], self.forward_kinematics(params, return_intermediates=True))
+        else:
+            return se2_action(start_positions[..., :], self.forward_kinematics(params, return_intermediates=False))
+
+    def _forward_kinematics_reduced(self, params, start_indices, end_indices):
+        # call check components before
+        current_action = np.zeros(params.shape[:-1] + (3,))
+        current_index = 0
+        for component in self.components:
+            is_turn = (start_indices <= current_index) & (current_index < end_indices)
+            component_action = component.forward_kinematics(params[..., current_index])
+            current_action[is_turn] = se2_compose(current_action[is_turn], component_action[is_turn])
+            if not isinstance(component, ConstComponent):
+                current_index += 1
+        return current_action # composed reduced action
 
 
 class Rail(RobotArm):
-    pass  # implement rail as components of const and simplerail
+    # implement rail as components of const and simplerail
+
+    def __init__(self, rail_angle):
+        before = ConstComponent(action=np.array([0, 0, rail_angle]))
+        rail_simple = SimpleRail()
+        after = ConstComponent(action=np.array([0, 0, -rail_angle]))
+        components = [before, rail_simple, after]
+        super().__init__(components)
 
 
 class Joint(RobotArm):
-    pass  # implement joint as components of const and simplejoint
+    # implement joint as components of const and simplejoint
+
+    def __init__(self, length, initial_angle):
+        before = ConstComponent(action=np.array([length, 0, initial_angle]))
+        joint_simple = SimpleJoint()
+        components = [before, joint_simple]
+        super().__init__(components)