From 388196c0e749231f725d5b02414298815547ce2b Mon Sep 17 00:00:00 2001 From: Lukas Date: Thu, 4 Jul 2024 03:02:56 +0200 Subject: [PATCH] implement robot_arm --- src/inverse_kinematics_sbi/components.py | 11 +++- src/inverse_kinematics_sbi/proba_robot.py | 2 +- src/inverse_kinematics_sbi/robots.py | 80 ++++++++++++++++++++--- 3 files changed, 80 insertions(+), 13 deletions(-) diff --git a/src/inverse_kinematics_sbi/components.py b/src/inverse_kinematics_sbi/components.py index 5cabd10..8501c1a 100644 --- a/src/inverse_kinematics_sbi/components.py +++ b/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 diff --git a/src/inverse_kinematics_sbi/proba_robot.py b/src/inverse_kinematics_sbi/proba_robot.py index c0b2092..12fb27d 100644 --- a/src/inverse_kinematics_sbi/proba_robot.py +++ b/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) diff --git a/src/inverse_kinematics_sbi/robots.py b/src/inverse_kinematics_sbi/robots.py index f6e27ce..e773ee8 100644 --- a/src/inverse_kinematics_sbi/robots.py +++ b/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)