Merge pull request #146 from cmower/add-sphere-colavoid

Add sphere collision avoidance constraints

example/sphere_collision_avoidance.py

@@ -0,0 +1,144 @@
+import os
+import pathlib
+
+import optas
+from optas.visualize import Visualizer
+
+def main(vis=True):
+
+    # Setup robot model
+    cwd = pathlib.Path(__file__).parent.resolve()  # path to current working directory
+    urdf_filename = os.path.join(cwd, "robots", "kuka_lwr", "kuka_lwr.urdf")
+    robot_model = optas.RobotModel(urdf_filename=urdf_filename, time_derivs=[0, 1, 2])
+    name = robot_model.get_name()
+    qnom = optas.deg2rad([0, -45, 0, 90, 0, 45, 0])
+    ee_link = "end_effector_ball"
+
+    T0 = robot_model.get_global_link_transform(ee_link, qnom)
+    z0 = T0[:3, 2]
+
+    def compute_initial_configuration():
+        # Compute an intiail configuration
+        start_eff_position = optas.DM([0.825, -0.35, 0.2])
+        builder = optas.OptimizationBuilder(1, robots=robot_model, derivs_align=True)
+        q = builder.get_model_state(name, 0)
+        T = robot_model.get_global_link_transform(ee_link, q)
+        p = T[:3, 3]
+        z = T[:3, 2]
+
+        builder.enforce_model_limits(name)
+
+        builder.add_equality_constraint("eff_pos", p, start_eff_position)
+        builder.add_equality_constraint("eff_ori", z, z0)
+
+        builder.initial_configuration(name, time_deriv=1)
+        builder.initial_configuration(name, time_deriv=2)
+
+        builder.add_cost_term("nominal", optas.sumsqr(q - qnom))
+
+        solver = optas.CasADiSolver(builder.build()).setup("ipopt")
+        solver.reset_initial_seed({f"{name}/q": qnom})
+        solution = solver.solve()
+
+        return solution[f"{name}/q"]
+
+
+    obs = optas.DM.zeros(3, 6)
+    obs[0, :] = 0.55
+    obs[1, :] = 0.0
+    obs[2, :] = optas.DM([0.1, 0.2, 0.3, 0.4, 0.5, 0.6]).T
+
+    obsrad = 0.1 * optas.DM.ones(6)
+
+    linkrad = 0.15
+
+    q0 = compute_initial_configuration()
+
+    duration = 10.
+
+    T = 20
+    builder = optas.OptimizationBuilder(T, robots=robot_model, derivs_align=True)
+
+    dt = duration / float(T - 1)
+
+    builder.enforce_model_limits(name)
+    builder.integrate_model_states(name, 2, dt)  # acc->vel
+    builder.integrate_model_states(name, 1, dt)  # vel->pos
+
+    builder.initial_configuration(name, init=q0)
+    builder.initial_configuration(name, time_deriv=1)
+    builder.initial_configuration(name, time_deriv=2)
+
+    goal_eff_position = optas.DM([0.825, 0.35, 0.2])
+    qf = builder.get_model_state(name, t=-1)
+    Tf = robot_model.get_global_link_transform(ee_link, qf)
+    pf = Tf[:3, 3]
+    zf = Tf[:3, 2]
+
+    builder.add_equality_constraint(
+        "eff_pos", pf, goal_eff_position, reduce_constraint=True
+    )
+    builder.add_equality_constraint("eff_ori", zf, z0, reduce_constraint=True)
+
+    builder.add_cost_term(
+        "min_vel", optas.sumsqr(builder.get_model_states(name, time_deriv=1))
+    )
+    builder.add_cost_term(
+        "min_acc", 100*optas.sumsqr(builder.get_model_states(name, time_deriv=2))
+    )
+
+    builder.add_cost_term("nominal", 1e2 * optas.sumsqr(qf - qnom))
+
+    obstacle_names = [f"obs{i}" for i in range(6)]
+
+
+    link_names = ['end_effector_ball', 'lwr_arm_7_link', 'lwr_arm_5_link', 'lwr_arm_6_link']
+
+    builder.sphere_collision_avoidance_constraints(name, obstacle_names, link_names=link_names)
+
+    solver = optas.CasADiSolver(builder.build()).setup("ipopt")
+
+    params = {}
+    for link_name in link_names:
+        params[link_name + "_radii"] = linkrad
+
+    for i, obstacle_name in enumerate(obstacle_names):
+        params[obstacle_name + "_position"] = obs[:, i]
+        params[obstacle_name + "_radii"] = obsrad[i]
+
+    solver.reset_parameters(params)
+
+    print("Solver ready")
+    solution = solver.solve()
+
+    if not vis:
+        return 0        
+
+    Q = solution[f"{name}/q"]
+
+    vis = Visualizer()
+    vis.robot_traj(robot_model, Q, animate=True, duration=duration)
+    vis.robot(robot_model, q0, alpha=0.5)
+    vis.grid_floor()
+    # vis.link(optas.DM.eye(4), axis_scale=1, axis_linewidth=8)
+    for i in range(6):
+        vis.sphere(radius=obsrad[i], position=obs[:, i], rgb=[1.0, 0.0, 0.0])
+
+    for link_name in link_names:
+        postraj = optas.DM.zeros(3, T)
+        for t in range(T):
+            postraj[:, t] = robot_model.get_global_link_position(link_name, Q[:, t])
+
+        vis.sphere_traj(
+            radius=linkrad,
+            position_traj=postraj,
+            rgb=[1, 0.64705882352, 0.0],
+            alpha_spec={'style': 'const', 'alpha': 0.25},
+            duration=duration,
+            animate=True,
+        )
+
+    vis.start()
+
+if __name__ == '__main__':
+    main()

optas/builder.py

@@ -336,17 +336,24 @@ def add_bound_inequality_constraint(
 
     @arrayify_args
     def add_equality_constraint(
-        self, name: str, lhs: CasADiArrayType, rhs: Union[None, CasADiArrayType] = None
+        self,
+        name: str,
+        lhs: CasADiArrayType,
+        rhs: Union[None, CasADiArrayType] = None,
+        reduce_constraint=False,
     ) -> None:
         """! Add the equality constraint lhs == rhs to the optimization problem.
 
         @param name Name for the constraint.
         @param lhs Left-hand side for the inequality constraint.
         @param rhs Right-hand side for the inequality constraint. If None, then it is replaced with the zero array with the same shape as lhs.
+        @param reduce_constraint When True, the constraints are specified by ||lhs - rhs||^2 = 0. This has the effect of reducing the number of equality constraints in the problem. However, if the constraints are linear they will be treated as nonlinear. Default is False.
         """
         if rhs is None:
             rhs = cs.DM.zeros(*lhs.shape)
         diff = rhs - lhs  # diff == 0
+        if reduce_constraint:
+            diff = cs.sumsqr(diff)
         if self._is_linear_in_x(diff):
             self._lin_eq_constraints[name] = diff
         else:
@@ -356,6 +363,52 @@ def add_equality_constraint(
     # Common constraints
     #
 
+    def sphere_collision_avoidance_constraints(
+        self, name, obstacle_names, link_names=None, base_link=None
+    ) -> None:
+        # Get model
+        model = self.get_model(name)
+        assert isinstance(model, RobotModel), "this method only applies to robot models"
+
+        # Update base link
+        if base_link is None:
+            base_link = model.get_root_link()
+
+        # Get model states
+        Q = self.get_model_states(name)
+        n = Q.shape[1]
+
+        # Setup inks
+        if link_names is None:
+            link_names = model.link_names
+        assert len(link_names), "at least one link should be named"
+        links = {}
+        for name in link_names:
+            links[name] = (
+                model.get_link_position_function(name, base_link),
+                self.add_parameter(name + "_radii"),
+            )
+
+        # Setup obstacles
+        assert len(obstacle_names), "at least one obstacle should be named"
+        obstacles = {}
+        for name in obstacle_names:
+            obstacles[name] = (
+                self.add_parameter(name + "_position", 3),
+                self.add_parameter(name + "_radii"),
+            )
+
+        # Add constraints
+        for t in range(n):
+            q = Q[:, t]
+            for link_name, (pos, linkrad) in links.items():
+                p = pos(q)
+                for obs_name, (obs, obsrad) in obstacles.items():
+                    name = f"sphere_col_avoid_{t}_{link_name}_{obs_name}"
+                    dist2 = cs.sumsqr(p - obs)
+                    bnd2 = (linkrad + obsrad) ** 2
+                    self.add_leq_inequality_constraint(name, bnd2, dist2)
+
     def integrate_model_states(
         self, name: str, time_deriv: int, dt: CasADiArrayType
     ) -> None:

optas/visualize.py

@@ -374,6 +374,8 @@ def sphere_traj(
             alpha_mid = alpha_spec.get("alpha_mid", 0.1)
             alpha_end = alpha_spec.get("alpha_end", 1.0)
             alphas = [alpha_start] + [alpha_mid] * (n - 2) + [alpha_end]
+        elif alpha_spec['style'] == 'const':
+            alphas = [alpha_spec.get("alpha", 1.0)]*n
 
         sphere_traj = []
         if animate:

tests/test_examples.py

@@ -89,6 +89,11 @@ def test_simple_joint_space_planner():
     assert main(gui=False) == 0
 
 
+def test_sphere_collision_avoidance():
+    main = load_main_function("sphere_collision_avoidance.py")
+    assert main(vis=False) == 0
+
+
 # def test_example():
 #     path = examples_path / "example.py"
 #     module = import_module(path)