Maximize Manipulability

maximize_logdet_Jacobian.py

import numpy as np
from pydrake.all import (
    AddMultibodyPlantSceneGraph,
    DiagramBuilder,
    Parser,
    MathematicalProgram,
    Solve,
    RigidTransform,
    JacobianWrtVariable,
    log,
    ExtractGradient,
    ExtractValue,
    AutoDiffXd,
)

def solve_manipulability_optimization():
    builder = DiagramBuilder()
    plant, _ = AddMultibodyPlantSceneGraph(builder, time_step=0.0)
    
    parser = Parser(plant)
    # Load the model
    iiwa_model = parser.AddModelsFromUrl(
        "package://drake_models/iiwa_description/urdf/iiwa14_primitive_collision.urdf")[0]
    
    # WELD THE BASE: Anchor the robot to the world origin
    plant.WeldFrames(
        plant.world_frame(), 
        plant.GetFrameByName("iiwa_link_0", iiwa_model)
    )
    
    plant.Finalize()
    
    # Now that it's welded, plant.num_positions() will correctly return 7
    num_q = plant.num_positions() 
    
    # Create AutoDiff plant for the cost function
    plant_ad = plant.ToAutoDiffXd()
    context_ad = plant_ad.CreateDefaultContext()
    
    frame_ad = plant_ad.GetFrameByName("iiwa_link_7")
    world_frame_ad = plant_ad.world_frame()

    prog = MathematicalProgram()
    q = prog.NewContinuousVariables(num_q, "q")

    def manipulability_cost(q_ad):
        plant_ad.SetPositions(context_ad, q_ad)
        
        # 1. Compute J (6x7)
        J_ad = plant_ad.CalcJacobianSpatialVelocity(
            context_ad, JacobianWrtVariable.kV, frame_ad,
            np.zeros(3), world_frame_ad, world_frame_ad
        )
        
        # 2. Extract values and gradients
        J_val = ExtractValue(J_ad) 
        dJdq_flat = ExtractGradient(J_ad) # (42 x 7)
        
        # 3. Regularize JJT for numerical health
        eps = 1e-6
        JJT = J_val @ J_val.T + eps * np.eye(6)
        
        # 4. Use slogdet (more stable for optimization)
        sign, logdet = np.linalg.slogdet(JJT)
        cost_val = -logdet
        
        # 5. Correct the Gradient Reshape
        # We use order='F' because Eigen/Drake is column-major
        dJdq = dJdq_flat.reshape((6, 7, num_q), order='F')
        
        inv_JJT = np.linalg.inv(JJT)
        cost_grad = np.zeros(num_q)
        
        for i in range(num_q):
            dJ_dqi = dJdq[:, :, i] 
            dJJT_dqi = dJ_dqi @ J_val.T + J_val @ dJ_dqi.T
            
            # Trace formula: d/dq -ln(det(A)) = -tr(A^-1 * dA/dq)
            cost_grad[i] = -np.trace(inv_JJT @ dJJT_dqi)

        print(np.round(ExtractValue(q_ad).flatten(), 2), "\t", cost_val)
        return AutoDiffXd(cost_val, cost_grad)

    prog.AddCost(manipulability_cost, q)
    prog.AddBoundingBoxConstraint(plant.GetPositionLowerLimits(), 
                                  plant.GetPositionUpperLimits(), q)

    # Initial guess
    np.random.seed(0)
    q0 = np.random.random(7) - 0.5

    # 1. Compute Initial Cost using the 'double' plant
    context = plant.CreateDefaultContext()
    plant.SetPositions(context, q0)
    
    J_init = plant.CalcJacobianSpatialVelocity(
        context, JacobianWrtVariable.kV, 
        plant.GetFrameByName("iiwa_link_7"),
        np.zeros(3), plant.world_frame(), plant.world_frame()
    )
    
    # Standard NumPy math on standard floats
    JJT_init = J_init @ J_init.T + 1e-6 * np.eye(6)
    initial_cost = -np.linalg.slogdet(JJT_init)[1]
    
    print(f"Initial Cost: {initial_cost:.4f}")

    # 2. Solve the optimization
    result = Solve(prog, q0)

    if result.is_success():
        final_cost = result.get_optimal_cost()
        print("Success!")
        print(f"Optimal q: {np.round(result.GetSolution(q), 3)}")
        print(f"Optimal Cost: {final_cost:.4f}")
        print(f"Improvement: {initial_cost - final_cost:.4f}")
    else:
        print(f"Failed: {result.get_solution_result()}")

if __name__ == "__main__":
    solve_manipulability_optimization()