Intro: KOMO - Motion Optimization

KOMO is a framework for designing motion by formulating constrained non-linear optimization problems. Inverse kinematics (IK) is the special case of optimizing only over a single configuration rather than a path. Formulating KOMO problems is key to realizing motion in rai.

The Features tutorial and the Notes: NLP Interface provide the mathematical background on differentiable features and the convention of how NLPs can be defined via features.

This tutorial shows how IK, rough waypoint optimization, and fine path optimization can be formulated as non-linear mathematical program (NLP) using KOMO. Essentially, the addObjective allows to add costs or constraints over any Feature to the NLP (same features that can be evaluated with C.eval).

Minimal IK example

Let’s first setup a single panda with box (located hanging in the air):

[ ]:

import robotic as ry
import numpy as np
import time

C = ry.Config()
C.addFile(ry.raiPath('scenarios/pandaSingle.g'))
C.view()

C.addFrame('box') \
    .setPosition([-.25,.1,1.]) \
    .setShape(ry.ST.ssBox, size=[.06,.06,.06,.005]) \
    .setColor([1,.5,0]) \
    .setContact(1)
C.view()

The following defines an optimization problem over a single configuration. The constructor declares over how many configurations (single, waypoints, path..) we optimize. The addObjective methods add costs or constraints:

[ ]:

qHome = C.getJointState()

komo = ry.KOMO(C, phases=1, slicesPerPhase=1, kOrder=0, enableCollisions=False)
komo.addObjective( # a sum-of-squares cost for staying close (in joint space) to qHome
    times=[],
    feature=ry.FS.jointState,
    frames=[],
    type=ry.OT.sos,
    scale=[1e-1],
    target=qHome
)
komo.addObjective([], ry.FS.positionDiff, ['l_gripper', 'box'], ry.OT.eq, [1e1]) # an equality constraint on the 3D position difference between gripper and box

We explain the KOMO constructor arguments later. (The above defines an IK problem.)

The addObjective method has signature

  • times: the time intervals (subset of configurations in a path) over which this feature is active (irrelevant for IK)

  • feature: the feature symbol (see the Features tutorial)

  • frames: the frames for which the feature is computed, given as list of frame names

  • type: whether this is a sum-of-squares (sos) cost, or eq or ineq constraint

  • scale: the matrix(!) by which the feature is multiplied

  • target: the offset which is substracted from the feature (before scaling)

Also concerning scale and target, see the Features tutorial) a mathematical description.

Given this definition of an optimization problem, we can call a generic NLP solver:

[ ]:

ret = ry.NLP_Solver(komo.nlp(), verbose=4) .solve()
print(ret)

With this high verbosity, individual newton steps and Augmented Lagrangian outer loops are displayed (we need only very few steps here).

The KOMO view displays the optimized configuration(s) stored by KOMO. (For paths, this is an overlay of many configurations. For IK, just one.):

[ ]:

komo.view(False, "IK solution")

We can get the sequence of joint state vectors for the optimized configuration(s) with getPath. Since this is only an IK problem, the sequence contains only the joint state vector for the single optimized configuration:

[ ]:

q = komo.getPath()
print(type(q), q.shape)

We’re done with KOMO and can destroy it. Then set the optimal joint state in C and view it:

[ ]:

del komo #also closes komo view
C.setJointState(q[0])
C.view()

Example for more constraints: box grasping IK

The key to design motions is to add clever constraints. Here is an example for more realistic box grasping:

[ ]:

komo = ry.KOMO(C, 1,1,0, True)
komo.addObjective([], ry.FS.jointState, [], ry.OT.sos, [1e-1], qHome)
komo.addObjective([], ry.FS.accumulatedCollisions, [], ry.OT.eq)
komo.addObjective([], ry.FS.jointLimits, [], ry.OT.ineq)
komo.addObjective([], ry.FS.positionDiff, ['l_gripper', 'box'], ry.OT.eq, [1e1])
komo.addObjective([], ry.FS.scalarProductXX, ['l_gripper', 'box'], ry.OT.eq, [1e1], [0])
komo.addObjective([], ry.FS.scalarProductXZ, ['l_gripper', 'box'], ry.OT.eq, [1e1], [0])
komo.addObjective([], ry.FS.distance, ['l_palm', 'box'], ry.OT.ineq, [1e1])

The two scalarProduct feature state that the gripper x-axis (which is the axis connecting the fingers) should be orthogonal to the object x- and z-axes. That implies fingers to normally oppose the object’s y-planes.

Note that grasping could also be opposing the object x- or z- planes – see below. Let solve it and then set the joint state to the solution:

[ ]:

ret = ry.NLP_Solver(komo.nlp(), verbose=0 ) .solve()
print(ret)
if ret.feasible:
    print('-- Always check feasibility flag of NLP solver return')
else:
    print('-- THIS IS INFEASIBLE!')

[ ]:

q = komo.getPath()
C.setJointState(q[0])
C.view(False, "IK solution")

Reusing the KOMO instance is ok if some aspect of the configuration changes and you want to resolve the same problem:

[ ]:

box = C.getFrame('box')
box.setPosition([-.25,.1,1.])
p0 = box.getPosition() # memory the start box position

for t in range(10):
    box.setPosition(p0 + .2 * np.random.randn(3)) # randomize box position
    komo.updateRootObjects(C) # only works for root objects (the 'box' is one)
    ret = ry.NLP_Solver(komo.nlp(), verbose=0 ) .solve()
    print(ret)
    q = komo.getPath()
    C.setJointState(q[0])
    C.view(False, 'IK solution - ' + ('*** INFEASIBLE ***' if not ret.feasible else 'feasible'))
    time.sleep(1.)

So the solver finds feasible grasps and exploits the null space of the constraints (grasps from different directions, but always opposing the y-planes).

To make this proper, we should actually test all three possible grasps - so let’s define 3 IK problems, solve each, and pick the best:

[ ]:

del komo
komo = []
for k in range(3):
    komo.append(ry.KOMO(C, 1,1,0, True))
    komo[k].addObjective([], ry.FS.jointState, [], ry.OT.sos, [1e-1], qHome)
    komo[k].addObjective([], ry.FS.accumulatedCollisions, [], ry.OT.eq)
    komo[k].addObjective([], ry.FS.jointLimits, [], ry.OT.ineq)
    komo[k].addObjective([], ry.FS.positionDiff, ['l_gripper', 'box'], ry.OT.eq, [1e1])
    komo[k].addObjective([], ry.FS.distance, ['l_palm', 'box'], ry.OT.ineq, [1e1])

komo[0].addObjective([], ry.FS.scalarProductXY, ['l_gripper', 'box'], ry.OT.eq, [1e1], [0])
komo[0].addObjective([], ry.FS.scalarProductXZ, ['l_gripper', 'box'], ry.OT.eq, [1e1], [0])

komo[1].addObjective([], ry.FS.scalarProductXX, ['l_gripper', 'box'], ry.OT.eq, [1e1], [0])
komo[1].addObjective([], ry.FS.scalarProductXZ, ['l_gripper', 'box'], ry.OT.eq, [1e1], [0])

komo[2].addObjective([], ry.FS.scalarProductXX, ['l_gripper', 'box'], ry.OT.eq, [1e1], [0])
komo[2].addObjective([], ry.FS.scalarProductXY, ['l_gripper', 'box'], ry.OT.eq, [1e1], [0])

[ ]:

for t in range(10):
    box.setPosition(p0 + .2 * np.random.randn(3))
    box.setQuaternion(np.random.randn(4)) # also set random orientation (quaternions get internally normalized)

    score = []
    for k in range(3):
        komo[k].updateRootObjects(C)
        ret = ry.NLP_Solver(komo[k].nlp(), verbose=0 ) .solve()
        score.append( 100.*(ret.eq+ret.ineq) + ret.sos )

    k = np.argmin(score)
    C.setJointState(komo[k].getPath()[0])
    C.view(False, f'IK solution {k} - ' + ('*** INFEASIBLE ***' if not ret.feasible else 'feasible'))
    time.sleep(1.)

[ ]:

del komo
del C

The above example is already quite elaborate: Defining multiple constraint problems, trying to solve them each, then picking the best. But this is good practise.

Waypoints example

Motion design can often be done by computing waypoints, i.e. a coarse-resolution sequence of poses. After this, you can either use the BotOp interface directly to interpolate them using splines, or invest more compute and compute optimal paths (fulfilling additional constraints) using path optimization between the waypoints, which we discuss later.

Let’s define a configuration where the desired gripper waypoints are pre-defined as marker frames. (That’s a common pattern: Simplify defining constraints by adding helper reference frames in the configuration.)

[ ]:

C = ry.Config()
C.addFile(ry.raiPath('scenarios/pandaSingle.g'))
C.addFrame('way1'). setShape(ry.ST.marker, [.1]) .setPosition([.4, .2, 1.])
C.addFrame('way2'). setShape(ry.ST.marker, [.1]) .setPosition([.4, .2, 1.4])
C.addFrame('way3'). setShape(ry.ST.marker, [.1]) .setPosition([-.4, .2, 1.])
C.addFrame('way4'). setShape(ry.ST.marker, [.1]) .setPosition([-.4, .2, 1.4])
qHome = C.getJointState()
C.view()

Now we can define a KOMO problem over 4 configurations, where at each configuration we impose position equality between gripper and a waypoint:

[ ]:

komo = ry.KOMO(C, phases=4, slicesPerPhase=1, kOrder=1, enableCollisions=False)
komo.addControlObjective([], 0, 1e-1)
komo.addControlObjective([], 1, 1e0)
komo.addObjective([1], ry.FS.positionDiff, ['l_gripper', 'way1'], ry.OT.eq, [1e1])
komo.addObjective([2], ry.FS.positionDiff, ['l_gripper', 'way2'], ry.OT.eq, [1e1])
komo.addObjective([3], ry.FS.positionDiff, ['l_gripper', 'way3'], ry.OT.eq, [1e1])
komo.addObjective([4], ry.FS.positionDiff, ['l_gripper', 'way4'], ry.OT.eq, [1e1])

ret = ry.NLP_Solver(komo.nlp(), verbose=0 ) .solve()
print(ret)
q = komo.getPath()
print(q)

for t in range(len(q)):
    C.setJointState(q[t])
    C.view(False, f'waypoint {t}')
    time.sleep(1)

The KOMO constructor has arguments:

  • config: the configuration, which is copied once (for IK) or many times (for waypoints/paths) to be the optimization variable

  • phases: the number \(P\) of phases (which essentially defines the real-valued interval \([0, P]\) over which objectives can be formulated)

  • slicesPerPhase: the discretizations per phase -> in total we have \(\texttt{phases} \cdot \texttt{slicesPerPhases}\) configurations which form the path and over which we optimize

  • kOrder: the “Markov-order”, i.e., maximal tuple of configurations over which we formulate features (e.g. take finite differences)

  • enableCollisions: if True, KOMO runs a broadphase collision check (using libFCL) in each optimization step – only then accumulative collision/penetration features will correctly evaluate to non-zero. But this is costly.

In our waypoint case: We have 4 phases, one for each waypoint. We don’t sub-sample the motion between waypoints, which is why we have slicesPerPhase=1. We formulate this as a 1-order problem: Some features take the finite difference between consecutive configurations (namely, to penalize velocities).

The addControlObjective is /almost/ the same as adding a FS.jointState objective: It penalizes distances in joint space. It has three arguments:

  • times: (as for addObjective) the phase-interval in which this objective holds; [] means all times

  • order: Do we penalize the jointState directly (order=0: penalizing sqr distance to qHome, order=1: penalizing sqr distances between consecutive configurations (velocities), order=2: penalizing accelerations across 3 configurations)

  • scale: as usual, but modulated by a factor “sqrt(delta t)” that somehow ensures total control costs in approximately independent of the choice of stepsPerPhase

In our waypoint case: We add control costs for both: homing (order 0, ensuring to stay close to homing), and velocities (order 1, penalizing movement between waypoints)

And the addObjective method now makes use of times argument: Specifying [1] means that this objective only holds in the interval [1,1], i.e. at phase-time 1 only.

Path example

Let’s do almost the same, but for a fine path. First order=1, leading to zig-zag, then order=2, leading to smooth path.

[ ]:

# Note, the slicesPerPhase=10 is the only difference to above
C.setJointState(qHome)
komo = ry.KOMO(C, 4, 10, 1, False)
komo.addControlObjective([], 0, 1e-1) # what happens if you change weighting to 1e0? why?
komo.addControlObjective([], 1, 1e0)
komo.addObjective([1], ry.FS.positionDiff, ['l_gripper', 'way1'], ry.OT.eq, [1e1])
komo.addObjective([2], ry.FS.positionDiff, ['l_gripper', 'way2'], ry.OT.eq, [1e1])
komo.addObjective([3], ry.FS.positionDiff, ['l_gripper', 'way3'], ry.OT.eq, [1e1])
komo.addObjective([4], ry.FS.positionDiff, ['l_gripper', 'way4'], ry.OT.eq, [1e1])

ret = ry.NLP_Solver(komo.nlp(), verbose=0 ) .solve()
print(ret)
q = komo.getPath()
print('size of path:', q.shape)

for t in range(q.shape[0]):
    C.setJointState(q[t])
    C.view(False, f'waypoint {t}')
    time.sleep(.1)

[ ]:

# only differences: the kOrder=2, control objective order 2, constrain final jointState velocity to zero
C.setJointState(qHome)
komo = ry.KOMO(C, 4, 10, 2, False)
komo.addControlObjective([], 0, 1e-1) # what happens if you change weighting to 1e0? why?
komo.addControlObjective([], 2, 1e0)
komo.addObjective([1], ry.FS.positionDiff, ['l_gripper', 'way1'], ry.OT.eq, [1e1])
komo.addObjective([2], ry.FS.positionDiff, ['l_gripper', 'way2'], ry.OT.eq, [1e1])
komo.addObjective([3], ry.FS.positionDiff, ['l_gripper', 'way3'], ry.OT.eq, [1e1])
komo.addObjective([4], ry.FS.positionDiff, ['l_gripper', 'way4'], ry.OT.eq, [1e1])
komo.addObjective([4], ry.FS.jointState, [], ry.OT.eq, [1e1], [], order=1)

ret = ry.NLP_Solver(komo.nlp(), verbose=0 ) .solve()
print(ret)
q = komo.getPath()
print('size of path:', q.shape)

for t in range(q.shape[0]):
    C.setJointState(q[t])
    C.view(False, f'waypoint {t}')
    time.sleep(.1)

Notice the new last objective! Without it, final velocity would not be zero. The last objective constrains the order=1 (i.e. velocity!) of the jointState feature to be zero.

Let’s plot the trajectory:

[ ]:

from matplotlib import pyplot as plt

print(q.shape)
fig, ax = plt.subplots()
ax.plot(q)
ax.set_xlabel("time slice (0=first-to-be-optimized)")
ax.set_ylabel("joint angles")
fig.tight_layout()
plt.show()

Let’s plot the errors/costs of the objectives over time:

[ ]:

err = komo.info_objectiveErrorTraces()
labels = komo.info_objectiveNames()

fig, ax = plt.subplots()
ax.plot(err, label=labels)
ax.legend(loc="upper left", prop={'size': 6})
ax.set_xlabel("time slice (0=first-to-be-optimized)")
ax.set_ylabel("objectiv error")
fig.tight_layout()
plt.show()

For fun, here is an interactive widget to slide through the solution. See the KOMO Viewer (raise to top). It follows the slider and displays the top error/cost objectives for the current time slice:

[ ]:

%matplotlib widget
from matplotlib.widgets import Slider

T = komo.getT()
komo.view()

fig, ax = plt.subplots()
fig.canvas.toolbar_visible = False
fig.canvas.header_visible = False
fig.canvas.footer_visible = False
fig.subplots_adjust(left=.1, bottom=.2, right=.95, top=.9)

ax.set_title('Use the slider to control view; raise KOMO Viewer "Always on Top"')
ax.plot(err, label=labels)
line, = ax.plot([0,0], ax.get_ylim(), linewidth=1.5, color='#b3b3b3')
ax.legend(loc="upper left", prop={'size': 6})
ax.set_xlabel("time slice (0=first-to-be-optimized)")
ax.set_ylabel("objective error")

def slider_callback(val):
    slice = max(-1, min(T-1, int(val)))
    if slice==-1:
        komo.view()
    else:
        komo.view_slice(slice, False)
        line.set_xdata([slice, slice])
    fig.canvas.draw_idle()

axslider = fig.add_axes([0.05, 0.05, 0.9, 0.03])
slider = Slider(ax=axslider, label='t', valmin=-1, valmax=T-1, valinit=0)
slider.on_changed(slider_callback)

plt.show()

[ ]:

plt.close('all')

[ ]:

del komo
del C

[ ]: