def prepare_ocp(
biorbd_model: biorbd.Model,
final_time: float,
n_shooting: int,
markers_ref: np.ndarray,
excitations_ref: np.ndarray,
q_ref: np.ndarray,
use_residual_torque: bool,
kin_data_to_track: str = "markers",
ode_solver: OdeSolver = OdeSolver.COLLOCATION(),
) -> OptimalControlProgram:
"""
Prepare the ocp to solve
Parameters
----------
biorbd_model: biorbd.Model
The loaded biorbd model
final_time: float
The time at final node
n_shooting: int
The number of shooting points
markers_ref: np.ndarray
The marker to track if 'markers' is chosen in kin_data_to_track
excitations_ref: np.ndarray
The muscle excitation (EMG) to track
q_ref: np.ndarray
The state to track if 'q' is chosen in kin_data_to_track
kin_data_to_track: str
The type of kin data to track ('markers' or 'q')
use_residual_torque: bool
If residual torque are present or not in the dynamics
ode_solver: OdeSolver
The ode solver to use
Returns
-------
The OptimalControlProgram ready to solve
"""
# Add objective functions
objective_functions = ObjectiveList()
objective_functions.add(ObjectiveFcn.Lagrange.TRACK_CONTROL,
key="muscles",
target=excitations_ref)
if use_residual_torque:
objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL,
key="tau")
if kin_data_to_track == "markers":
objective_functions.add(ObjectiveFcn.Lagrange.TRACK_MARKERS,
node=Node.ALL,
weight=100,
target=markers_ref)
elif kin_data_to_track == "q":
objective_functions.add(
ObjectiveFcn.Lagrange.TRACK_STATE,
key="q",
weight=100,
node=Node.ALL,
target=q_ref,
index=range(biorbd_model.nbQ()),
)
else:
raise RuntimeError("Wrong choice of kin_data_to_track")
# Dynamics
dynamics = DynamicsList()
dynamics.add(DynamicsFcn.MUSCLE_DRIVEN,
with_excitations=True,
with_residual_torque=use_residual_torque)
# Path constraint
x_bounds = BoundsList()
x_bounds.add(bounds=QAndQDotBounds(biorbd_model))
# Due to unpredictable movement of the forward dynamics that generated the movement, the bound must be larger
x_bounds[0].min[[0, 1], :] = -2 * np.pi
x_bounds[0].max[[0, 1], :] = 2 * np.pi
# Add muscle to the bounds
activation_min, activation_max, activation_init = 0, 1, 0.5
x_bounds[0].concatenate(
Bounds([activation_min] * biorbd_model.nbMuscles(),
[activation_max] * biorbd_model.nbMuscles()))
x_bounds[0][(biorbd_model.nbQ() + biorbd_model.nbQdot()):,
0] = excitations_ref[:, 0]
# Initial guess
x_init = InitialGuessList()
x_init.add([0] * (biorbd_model.nbQ() + biorbd_model.nbQdot()) +
[0] * biorbd_model.nbMuscles())
# Define control path constraint
excitation_min, excitation_max, excitation_init = 0, 1, 0.5
u_bounds = BoundsList()
u_init = InitialGuessList()
if use_residual_torque:
tau_min, tau_max, tau_init = -100, 100, 0
u_bounds.add(
[tau_min] * biorbd_model.nbGeneralizedTorque() +
[excitation_min] * biorbd_model.nbMuscles(),
[tau_max] * biorbd_model.nbGeneralizedTorque() +
[excitation_max] * biorbd_model.nbMuscles(),
)
u_init.add([tau_init] * biorbd_model.nbGeneralizedTorque() +
[excitation_init] * biorbd_model.nbMuscles())
else:
u_bounds.add([excitation_min] * biorbd_model.nbMuscles(),
[excitation_max] * biorbd_model.nbMuscles())
u_init.add([excitation_init] * biorbd_model.nbMuscles())
# ------------- #
return OptimalControlProgram(
biorbd_model,
dynamics,
n_shooting,
final_time,
x_init,
u_init,
x_bounds,
u_bounds,
objective_functions,
ode_solver=ode_solver,
)
def generate_data(biorbd_model: biorbd.Model,
final_time: float,
n_shooting: int,
use_residual_torque: bool = True) -> tuple:
"""
Generate random data. If np.random.seed is defined before, it will always return the same results
Parameters
----------
biorbd_model: biorbd.Model
The loaded biorbd model
final_time: float
The time at final node
n_shooting: int
The number of shooting points
use_residual_torque: bool
If residual torque are present or not in the dynamics
Returns
-------
The time, marker, states and controls of the program. The ocp will try to track these
"""
# Aliases
n_q = biorbd_model.nbQ()
n_qdot = biorbd_model.nbQdot()
n_tau = biorbd_model.nbGeneralizedTorque()
n_mus = biorbd_model.nbMuscleTotal()
dt = final_time / n_shooting
# Casadi related stuff
symbolic_q = MX.sym("q", n_q, 1)
symbolic_qdot = MX.sym("qdot", n_qdot, 1)
symbolic_mus_states = MX.sym("mus", n_mus, 1)
symbolic_tau = MX.sym("tau", n_tau, 1)
symbolic_mus_controls = MX.sym("mus", n_mus, 1)
symbolic_states = vertcat(*(symbolic_q, symbolic_qdot,
symbolic_mus_states))
symbolic_controls = vertcat(*(symbolic_tau, symbolic_mus_controls))
symbolic_parameters = MX.sym("u", 0, 0)
nlp = NonLinearProgram()
nlp.model = biorbd_model
nlp.variable_mappings = {
"q": BiMapping(range(n_q), range(n_q)),
"qdot": BiMapping(range(n_qdot), range(n_qdot)),
"tau": BiMapping(range(n_tau), range(n_tau)),
"muscles": BiMapping(range(n_mus), range(n_mus)),
}
markers_func = biorbd.to_casadi_func("ForwardKin", biorbd_model.markers,
symbolic_q)
nlp.states.append("q", [symbolic_q, symbolic_q], symbolic_q,
nlp.variable_mappings["q"])
nlp.states.append("qdot", [symbolic_qdot, symbolic_qdot], symbolic_qdot,
nlp.variable_mappings["qdot"])
nlp.states.append("muscles", [symbolic_mus_states, symbolic_mus_states],
symbolic_mus_states, nlp.variable_mappings["muscles"])
nlp.controls.append("tau", [symbolic_tau, symbolic_tau], symbolic_tau,
nlp.variable_mappings["tau"])
nlp.controls.append(
"muscles",
[symbolic_mus_controls, symbolic_mus_controls],
symbolic_mus_controls,
nlp.variable_mappings["muscles"],
)
dynamics_func = biorbd.to_casadi_func(
"ForwardDyn",
DynamicsFunctions.muscles_driven,
symbolic_states,
symbolic_controls,
symbolic_parameters,
nlp,
False,
)
def dyn_interface(t, x, u):
u = np.concatenate([np.zeros(n_tau), u])
return np.array(dynamics_func(x, u, np.empty((0, 0)))).squeeze()
# Generate some muscle excitations
U = np.random.rand(n_shooting, n_mus).T
# Integrate and collect the position of the markers accordingly
X = np.ndarray((n_q + n_qdot + n_mus, n_shooting + 1))
markers = np.ndarray((3, biorbd_model.nbMarkers(), n_shooting + 1))
def add_to_data(i, q):
X[:, i] = q
markers[:, :, i] = markers_func(q[:n_q])
x_init = np.array([0] * n_q + [0] * n_qdot + [0.5] * n_mus)
add_to_data(0, x_init)
for i, u in enumerate(U.T):
sol = solve_ivp(dyn_interface, (0, dt),
x_init,
method="RK45",
args=(u, ))
x_init = sol["y"][:, -1]
add_to_data(i + 1, x_init)
time_interp = np.linspace(0, final_time, n_shooting + 1)
return time_interp, markers, X, U
def prepare_ocp(
biorbd_model: biorbd.Model,
final_time: float,
n_shooting: int,
markers_ref: np.ndarray,
tau_ref: np.ndarray,
ode_solver: OdeSolver = OdeSolver.RK4(),
) -> OptimalControlProgram:
"""
Prepare the ocp
Parameters
----------
biorbd_model: biorbd.Model
The loaded biorbd model
final_time: float
The time at final node
n_shooting: int
The number of shooting points
markers_ref: np.ndarray
The markers to track
tau_ref: np.ndarray
The generalized forces to track
ode_solver: OdeSolver
The ode solver to use
Returns
-------
The OptimalControlProgram ready to be solved
"""
# Add objective functions
objective_functions = ObjectiveList()
objective_functions.add(
ObjectiveFcn.Lagrange.TRACK_MARKERS,
axes=[Axis.Y, Axis.Z],
node=Node.ALL,
weight=100,
target=markers_ref[1:, :, :],
)
objective_functions.add(ObjectiveFcn.Lagrange.TRACK_CONTROL,
key="tau",
target=tau_ref)
# Dynamics
dynamics = DynamicsList()
expand = False if isinstance(ode_solver, OdeSolver.IRK) else True
dynamics.add(DynamicsFcn.TORQUE_DRIVEN, expand=expand)
# Path constraint
x_bounds = BoundsList()
x_bounds.add(bounds=QAndQDotBounds(biorbd_model))
x_bounds[0][:, 0] = 0
# Initial guess
n_q = biorbd_model.nbQ()
n_qdot = biorbd_model.nbQdot()
x_init = InitialGuessList()
x_init.add([0] * (n_q + n_qdot))
# Define control path constraint
n_tau = biorbd_model.nbGeneralizedTorque()
tau_min, tau_max, tau_init = -100, 100, 0
u_bounds = BoundsList()
u_bounds.add([tau_min] * n_tau, [tau_max] * n_tau)
u_init = InitialGuessList()
u_init.add([tau_init] * n_tau)
# ------------- #
return OptimalControlProgram(
biorbd_model,
dynamics,
n_shooting,
final_time,
x_init,
u_init,
x_bounds,
u_bounds,
objective_functions,
ode_solver=ode_solver,
)