Skip to content

Commit

Permalink
Update model reduction logic
Browse files Browse the repository at this point in the history
  • Loading branch information
@flferretti
flferretti committed Jan 19, 2024
1 parent 874c16c commit 01533b6
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 68 deletions.
54 changes: 18 additions & 36 deletions src/adam/core/rbd_algorithms.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -481,15 +481,7 @@ def aba(
joint_accelerations (T): The joints acceleration
"""
model = self.model.reduce(self.model.actuated_joints)

joints = list(
filter(
lambda joint: joint.name in self.model.actuated_joints,
self.model.joints.values(),
)
)

joints.sort(key=lambda joint: joint.idx)
joints = list(model.joints.values())

NB = model.N

Expand All @@ -506,10 +498,8 @@ def aba(
sdd = self.math.factory.zeros(NB, 1, 1)
B_X_W = self.math.adjoint_mixed(base_transform)

if self.model.floating_base:
IA[0] = self.model.tree.get_node_from_name(
self.root_link
).link.spatial_inertia()
if model.floating_base:
IA[0] = model.tree.get_node_from_name(self.root_link).link.spatial_inertia()
v[0] = B_X_W @ base_velocity
pA[0] = (
self.math.spatial_skew_star(v[0]) @ IA[0] @ v[0]
Expand All @@ -522,7 +512,7 @@ def get_tree_transform(self, joints) -> "Array":
Array: the tree transform
"""
relative_transform = lambda j: self.math.inv(
self.model.tree.graph[j.parent].parent_arc.spatial_transform(0)
model.tree.graph[j.child].parent_arc.spatial_transform(0)
) @ j.spatial_transform(0)

return self.math.vertcat(
Expand All @@ -538,19 +528,7 @@ def get_tree_transform(self, joints) -> "Array":
)

tree_transform = get_tree_transform(self, joints)

find_parent = (
lambda j: find_parent(model.tree.get_node_from_name(j.parent).parent_arc)
if model.tree.get_node_from_name(j.parent).parent_arc.idx is None
else model.tree.get_node_from_name(j.parent).parent_arc.idx
)

p = [-1] + [
model.tree.get_idx_from_name(i.parent)
if model.tree.get_idx_from_name(i.parent) < NB
else find_parent(i)
for i in joints
]
p = lambda i: list(model.tree.graph).index(joints[i].parent)

# Pass 1
for i, joint in enumerate(joints[1:], start=1):
Expand All @@ -561,8 +539,8 @@ def get_tree_transform(self, joints) -> "Array":
i_X_pi[i] = joint.spatial_transform(q) @ tree_transform[i]
v_J = joint.motion_subspace() * q_dot

v[i] = i_X_pi[i] @ v[p[i]] + v_J
c[i] = i_X_pi[i] @ c[p[i]] + self.math.spatial_skew(v[i]) @ v_J
v[i] = i_X_pi[i] @ v[p(i)] + v_J
c[i] = i_X_pi[i] @ c[p(i)] + self.math.spatial_skew(v[i]) @ v_J

IA[i] = model.tree.get_node_from_name(joint.parent).link.spatial_inertia()

Expand All @@ -579,17 +557,21 @@ def get_tree_transform(self, joints) -> "Array":
):
U[i] = IA[i] @ joint.motion_subspace()
D[i] = joint.motion_subspace().T @ U[i]
u[i] = self.math.vertcat(tau[joint.idx]) - joint.motion_subspace().T @ pA[i]
u[i] = (
self.math.vertcat(tau[joint.idx]) - joint.motion_subspace().T @ pA[i]
if joint.idx is not None
else 0.0
)

Ia = IA[i] - U[i] / D[i] @ U[i].T
pa = pA[i] + Ia @ c[i] + U[i] * u[i] / D[i]

if joint.parent != self.root_link or not self.model.floating_base:
IA[p[i]] += i_X_pi[i].T @ Ia @ i_X_pi[i]
pA[p[i]] += i_X_pi[i].T @ pa
if joint.parent != self.root_link or not model.floating_base:
IA[p(i)] += i_X_pi[i].T @ Ia @ i_X_pi[i]
pA[p(i)] += i_X_pi[i].T @ pa
continue

a[0] = B_X_W @ g if self.model.floating_base else self.math.solve(-IA[0], pA[0])
a[0] = B_X_W @ g if model.floating_base else self.math.solve(-IA[0], pA[0])

# Pass 3
for i, joint in enumerate(joints[1:], start=1):
Expand All @@ -598,7 +580,7 @@ def get_tree_transform(self, joints) -> "Array":

sdd[i - 1] = (u[i] - U[i].T @ a[i]) / D[i]

a[i] += i_X_pi[i].T @ a[p[i]] + joint.motion_subspace() * sdd[i - 1] + c[i]
a[i] += i_X_pi[i].T @ a[p(i)] + joint.motion_subspace() * sdd[i - 1] + c[i]

# Squeeze sdd
s_ddot = self.math.vertcat(*[sdd[i] for i in range(sdd.shape[0])])
Expand All @@ -613,7 +595,7 @@ def get_tree_transform(self, joints) -> "Array":
return self.math.horzcat(
self.math.vertcat(
self.math.solve(B_X_W, a[0]) + g
if self.model.floating_base
if model.floating_base
else self.math.zeros(6, 1),
),
s_ddot,
Expand Down
12 changes: 5 additions & 7 deletions src/adam/model/model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -98,14 +98,12 @@ def reduce(self, joints_name_list: List[str]) -> "Model":
)

tree = self.tree.reduce(joints_name_list)
joints_list = list(
filter(
lambda joint: joint.name in self.actuated_joints,
self.joints.values(),
)
)

joints_list = [
node.parent_arc for node in tree.graph.values() if node.name != tree.root
]
joints_list.sort(key=lambda joint: joint.idx)
# update nodes dict

links = {node.name: node.link for node in tree.graph.values()}
joints = {joint.name: joint for joint in joints_list}
frames = {
Expand Down
44 changes: 19 additions & 25 deletions src/adam/model/tree.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,15 +82,6 @@ def reduce(self, considered_joint_names: List[str]) -> "Tree":
Returns:
Tree: the reduced tree
"""
# find the nodes between two fixed joints
nodes_to_lump = list(
{
joint.child
for node in self.graph.values()
for joint in node.arcs
if joint.name not in considered_joint_names
}
)

relative_transform = (
lambda node: node.link.math.inv(
Expand All @@ -101,21 +92,17 @@ def reduce(self, considered_joint_names: List[str]) -> "Tree":
else node.parent_arc.spatial_transform(0)
)

last = []
leaves = [node for node in self.graph.values() if node.children == last]
# find the tree leaves and proceed until the root
leaves = [node for node in self.graph.values() if node.children == []]

while all(leaf.name != self.root for leaf in leaves):
for leaf in leaves:
if leaf is self.graph[self.root]:
continue

if leaf.parent_arc.name not in considered_joint_names:
# create the new node
if leaf.parent_arc.name not in considered_joint_names + [self.root]:
new_node = Node(
name=leaf.parent.name,
link=None,
arcs=[],
children=None,
children=[],
parent=None,
parent_arc=None,
)
Expand All @@ -129,22 +116,29 @@ def reduce(self, considered_joint_names: List[str]) -> "Tree":
# update the parents
new_node.parent = self.graph[leaf.parent.name].parent
new_node.parent_arc = self.graph[new_node.name].parent_arc
new_node.parent_arc.parent = (
leaf.children[0].parent_arc.name if leaf.children != [] else []
)

# update the children
new_node.children = leaf.children
new_node.children = [
child for child in leaf.children if child.name in self.graph
]

for child in new_node.children:
child.parent = new_node.link
child.parent_arc = new_node.parent_arc

# update the arcs
if leaf.arcs != []:
for arc in leaf.arcs:
if arc.name in considered_joint_names:
new_node.arcs.append(arc)
new_node.arcs = (
[arc for arc in leaf.arcs if arc.name in considered_joint_names]
if leaf.arcs != []
else []
)
for j in new_node.arcs:
j.parent = new_node.link.name

logging.debug(f"Removing {leaf.name}")
self.graph.pop(leaf.name)
self.graph[new_node.name] = new_node
self.ordered_nodes_list.remove(leaf.name)
leaves = [
self.get_node_from_name((leaf.parent.name))
for leaf in leaves
Expand Down

0 comments on commit 01533b6

Please sign in to comment.