pendulum_regrasp.py

from pendulum_croco_models import (
    ActuationModelDoublePendulum,
    CostModelDoublePendulum,
)
from utils import get_args, computeStickingCondition
from multibody_with_lugre_state import MultiBodyWithLugreState

import matplotlib.pyplot as plt
import crocoddyl
import signal
import time
import example_robot_data
import numpy as np

args = get_args()

WITHDISPLAY = args.visualizer
WITHPLOT = args.plotter
signal.signal(signal.SIGINT, signal.SIG_DFL)


# Loading the double pendulum model
pendulum = example_robot_data.load("double_pendulum")
pendulum.model.effortLimit[0] = 10
pendulum.model.effortLimit[1] = 10
model = pendulum.model
data = pendulum.data

state = crocoddyl.StateMultibody(model)
if args.friction_model == "lugre":
    state = MultiBodyWithLugreState(state)
# this is the first link
actuation = ActuationModelDoublePendulum(state, args, pendulum.model, pendulum.data)

nu = actuation.nu
runningCostModel = crocoddyl.CostModelSum(state, nu)
terminalCostModel = crocoddyl.CostModelSum(state, nu)

xResidual = crocoddyl.ResidualModelState(state, state.zero(), nu)
xActivation = crocoddyl.ActivationModelQuad(state.ndx)
uResidual = crocoddyl.ResidualModelControl(state, nu)
xRegCost = crocoddyl.CostModelResidual(state, xActivation, xResidual)
uRegCost = crocoddyl.CostModelResidual(state, uResidual)
xPendCost = CostModelDoublePendulum(
    state, crocoddyl.ActivationModelWeightedQuad(np.array([1.0] * 4 + [0.1] * 2)), nu
)

dt = 1e-3

runningCostModel.addCost("uReg", uRegCost, 1e-4 / dt)
runningCostModel.addCost("xGoal", xPendCost, 1e-5 / dt)
terminalCostModel.addCost("xGoal", xPendCost, 100.0)

# NOTE: this can't take anything but multibodystate
# TODO: just start implementing all this shit in c++ already,
# we've seen enough python and source code to know what's going on
# (or at least it won't be any better with more prep...)
runningModel = crocoddyl.IntegratedActionModelEuler(
    crocoddyl.DifferentialActionModelFreeFwdDynamics(
        state, actuation, runningCostModel
    ),
    dt,
)
runningModel.u_lb = np.array([-20.0, 0.10])
runningModel.u_ub = np.array([20.0, 100.0])
terminalModel = crocoddyl.IntegratedActionModelEuler(
    crocoddyl.DifferentialActionModelFreeFwdDynamics(
        state, actuation, terminalCostModel
    ),
    dt,
)
terminalModel.u_lb = np.array([-20.0, 0.10])
terminalModel.u_ub = np.array([20.0, 100.0])

# Creating the shooting problem and the FDDP solver
# TODO: WHY IN THE NAME OF CHRIST DOES CHANGING T CHANGE x0???????????????
# ???????????????????????????????????????????????????????????????????????
# ???????????????????????????????????????????????????????????????????????
# ???????????????????????????????????????????????????????????????????????
T = 500
x0 = np.array([5 * np.pi / 6, 0.0, 0.0, 0.0])
problem = crocoddyl.ShootingProblem(x0, [runningModel] * T, terminalModel)
# solver = crocoddyl.SolverFDDP(problem)
solver = crocoddyl.SolverBoxFDDP(problem)
if WITHPLOT:
    solver.setCallbacks(
        [
            crocoddyl.CallbackVerbose(),
            crocoddyl.CallbackLogger(),
        ]
    )
else:
    solver.setCallbacks([crocoddyl.CallbackVerbose()])

# Solving the problem with the FDDP solver
solver.solve()
# Plotting the entire motion
log = solver.getCallbacks()[1]
print(x0)
xs = np.array(log.xs)
us = np.array(log.us)
print(xs[0])
if WITHPLOT:
    crocoddyl.plotOCSolution(log.xs, log.us, figIndex=1, show=False)
    crocoddyl.plotConvergence(
        log.costs, log.pregs, log.dregs, log.grads, log.stops, log.steps, figIndex=2
    )

# TODO: plot sticking and sliding conditions
sticking_conditions = []
for x, u in zip(xs, us):
    sticking_conditions.append(
        computeStickingCondition(actuation.mu_tors, pendulum.model, pendulum.data, x, u)
    )
plt.plot(np.arange(len(sticking_conditions)), sticking_conditions)
plt.title("sticking is 1, sliding is 0")
plt.show()

# Display the entire motion
if WITHDISPLAY:
    display = crocoddyl.MeshcatDisplay(pendulum)
    display.rate = -1
    display.freq = 1
    while True:
        display.displayFromSolver(solver)
        time.sleep(1.0)