in differentiable_robot_model/robot_model.py [0:0]
def update_kinematic_state(self, q: torch.Tensor, qd: torch.Tensor) -> None:
r"""
Updates the kinematic state of the robot
Args:
q: joint angles [batch_size x n_dofs]
qd: joint velocities [batch_size x n_dofs]
Returns:
"""
assert q.ndim == 2
assert qd.ndim == 2
assert q.shape[1] == self._n_dofs
assert qd.shape[1] == self._n_dofs
batch_size = q.shape[0]
# update the state of the joints
for i in range(q.shape[1]):
idx = self._controlled_joints[i]
self._bodies[idx].update_joint_state(
q[:, i].unsqueeze(1), qd[:, i].unsqueeze(1)
)
# we assume a non-moving base
parent_body = self._bodies[0]
parent_body.vel = SpatialMotionVec(
torch.zeros((batch_size, 3), device=self._device),
torch.zeros((batch_size, 3), device=self._device),
)
# propagate the new joint state through the kinematic chain to update bodies position/velocities
for i in range(1, len(self._bodies)):
body = self._bodies[i]
parent_name = self._urdf_model.get_name_of_parent_body(body.name)
# find the joint that has this link as child
parent_body = self._bodies[self._name_to_idx_map[parent_name]]
# transformation operator from child link to parent link
childToParentT = body.joint_pose
# transformation operator from parent link to child link
parentToChildT = childToParentT.inverse()
# the position and orientation of the body in world coordinates, with origin at the joint
body.pose = parent_body.pose.multiply_transform(childToParentT)
# we rotate the velocity of the parent's body into the child frame
new_vel = parent_body.vel.transform(parentToChildT)
# this body's angular velocity is combination of the velocity experienced at it's parent's link
# + the velocity created by this body's joint
body.vel = body.joint_vel.add_motion_vec(new_vel)
return