n_phases = 2
import robotic as ry
import numpy as np

komo = ry.KOMO(config, n_phases, 10, 2, True)
# komo.addObjective([], ry.FS.accumulatedCollisions, [], ry.OT.eq)
# komo.addObjective([], ry.FS.jointLimits, [], ry.OT.ineq)
komo.addControlObjective([], 0, 1e-1)
komo.addControlObjective([], 1, 1e0)

# Phase 1: Align gripper with big red block
komo.addObjective([0.], ry.FS.positionDiff, ['l_gripper', 'big_red_block'], ry.OT.eq, [1e2], [0, 0, 0.1])
komo.addObjective([0.], ry.FS.vectorZ, ['l_gripper'], ry.OT.eq, [1e2], [0, 0, 1])
komo.addObjective([0.], ry.FS.scalarProductXX, ['l_gripper', 'big_red_block'], ry.OT.eq, [1e2], [1])

# Phase 2: Push towards target pose
komo.addObjective([1.], ry.FS.positionDiff, ['big_red_block', 'target_pose'], ry.OT.eq, [1e2])
komo.addObjective([1.], ry.FS.poseDiff, ['big_red_block', 'target_pose'], ry.OT.eq, [1e2])

# Avoid collision with table during push
komo.addObjective([0., 1.], ry.FS.distance, ['l_gripper', 'table'], ry.OT.ineq, [1e2], [0.05])
