Go2 Walking (MPC)¶
examples/mujoco_mpc/go2/go2_walking.py
walks a Unitree Go2 quadruped around in a live MuJoCo simulation, with a trot driven by a
gamepad: the sticks command a forward, sideways and turning velocity, and further buttons change the base height,
tilt and step height, or make the robot jump in place. Everything is planned by a single whole-body
MPC that solves for joint accelerations and contact forces, and the torques that come out of it
are applied to the simulator directly.
It combines two earlier examples. From Go2 Jump it takes the Go2 model, the
ReducedInverseDynamics solver mode and the idea of scheduling contacts per
stage, but where the jump fixed its schedule once, here it is regenerated at every plan stage from the joystick.
From Panda Circle (MPC) it takes the receding-horizon loop and the pattern of keeping the
costs in lists, so that each step can point them at a different target. What is new is the part in the middle: the
code that turns "go forward at 0.4 m/s" into a base trajectory, a set of footholds and swing paths, and finally a
setting for each cost and constraint of the model.
The model itself knows nothing about gaits. This is worth keeping in mind while reading the script: the SSQP model is built once, as a list of generic costs and contact constraints, and everything that makes it walk is in how they are updated.
Timing¶
description="Go2 walking MPC in MuJoCo using SSQP"
)
parser.add_argument(
"-q",
"--qp-solver",
choices=["admm", "hpipm", "piqp"],
default="hpipm",
help="QP solver used at each SSQP iteration",
)
return parser.parse_args()
There are two clocks. MuJoCo runs at dt = 2 ms. The MPC has 20 stages of mpc_dt = 24 ms, so it plans 0.48 s
ahead. They are related by STAGE_EVERY, the 12 simulation steps in one plan stage, and mpc_dt is required to be a
multiple of dt for this to be a whole number. The script keeps a counter stage that advances once per plan stage,
not once per simulation step, and the whole planner is written in stages: the gait period, the jump phases and the
base references are all counted in them.
The MPC is also solved once per stage, which is every 24 ms of simulation time. Solving every simulation step would
be 12 times more work for little gain, since the plan only has a new first stage every mpc_dt anyway. The simulation
is synced with the controller, so if a solve takes longer than 24 ms of real time, the simulation simply runs slower
than real time.
Weights¶
ssqpy.setSilentMode()
np.set_printoptions(suppress=True)
dt = 0.002
mpc_H = 20
# Keep mpc_dt a multiple of dt, so that plan stages line up with sim steps
mpc_dt = 0.024
# stage indexes the plan (mpc_dt per step), so only advance it once per
# planning step, not every simulation step
STAGE_EVERY = round(mpc_dt / dt)
# Print and sync the viewer once per plan stage, so that each measured window
# holds exactly one MPC solve
T_WGT and R_WGT weigh position and rotation tracking, the base's and the feet's. C_WGT is a weak pull of the
joints toward a nominal standing configuration, V_WGT penalizes joint velocity, and U_WGT and F_WGT penalize the
control and the contact forces. The meaning of the control depends on the dynamics mode: in forward dynamics it is a
joint torque, but in ReducedInverseDynamics it is a joint acceleration, which has very different units, so the
weight is different for each (with U_WGT much smaller for accelerations). Changing the mode in this script without
checking the weights would silently retune the controller.
Building the model¶
JUMP_LAND_STAGES = 12
# Base height at takeoff and touchdown
JUMP_TAKEOFF_HEIGHT = 0.36
# Weight holding the legs in their takeoff configuration while in the air,
# where the feet and the base's path can no longer be steered
JUMP_CONFIG_WGT = 1.0
GRAVITY = 9.81
JUMP_SPEED = GRAVITY * JUMP_FLIGHT_STAGES * mpc_dt / 2
JUMP_CROUCH_HEIGHT = (
JUMP_TAKEOFF_HEIGHT - JUMP_SPEED * JUMP_PUSH_STAGES * mpc_dt / 2
)
JUMP_LAND_HEIGHT = (
JUMP_TAKEOFF_HEIGHT - JUMP_SPEED * JUMP_LAND_STAGES * mpc_dt / 2
)
THIGH_LENGTH = 0.213
# Distance from the knee to the foot frame
CALF_LENGTH = 0.23648
model = ssqpy.model.Model(
mpc_H,
mpc_dt,
mjcf_path="../../assets/mjcf/go2/go2.xml",
free_flyer=True,
solver_mode=dynamics_mode,
)
nq = model.getnq()
nv = model.getnv()
nu = model.getnu()
The model is built from the robot's MJCF file with free_flyer=True (the Go2 Jump example uses the
Go2 preset instead). MuJoCo simulates the same robot description, in a scene file that adds the floor and
obstacles. The initial configuration places the base at BASE_HEIGHT of
0.34 m and sets every leg with leg_config, the closed-form hip, thigh and calf angles that put the foot straight
below the hip at that height (a two-link inverse kinematics problem, using the thigh and calf lengths). The base
quaternion is the identity, so the robot starts facing along \(+x\), away from the obstacles in the scene.
# Hip, thigh and calf angles placing the foot straight below the hip
calf = -np.arccos(
(height**2 - THIGH_LENGTH**2 - CALF_LENGTH**2)
/ (2 * THIGH_LENGTH * CALF_LENGTH)
)
thigh = np.arcsin(-CALF_LENGTH * np.sin(calf) / height)
return np.array([0.0, thigh, calf])
base_pos = np.array([0.0, 0.0, BASE_HEIGHT])
base_config = np.tile(leg_config(BASE_HEIGHT), 4)
# Initial configuration
init_state = np.zeros(37)
init_state[:3] = base_pos
init_state[6] = 1.0
init_state[7:19] = base_config
foot_names = ["FL_foot", "FR_foot", "RL_foot", "RR_foot"]
# Trot pairs, swinging in alternate gait periods
SWING_PAIRS = [["FL_foot", "RR_foot"], ["FR_foot", "RL_foot"]]
# The base starts facing along the world x axis, away from the scene's
# obstacles
R = np.eye(3)
# Foot positions relative to the base, on the ground
Two sets of geometry are computed from the initial pose with forward kinematics. foot_offsets is each foot's
position relative to the base, flattened to the ground, which is where the foot would stand when the robot is
upright; the planner uses it as the foot's nominal position under the base. hip_offsets and HIP_LENGTH give the
hip geometry that the leg inverse kinematics (leg_ik, further down in the script) needs, since the Go2's thigh is
offset sideways from the hip joint.
The gait¶
A trot moves diagonal pairs of feet together. SWING_PAIRS lists the two pairs, and they take turns: in even
gait periods FL and RR swing while FR and RL stand, and in odd periods the roles swap. GAIT_PERIOD is 10 stages
(0.24 s), so every foot swings for 0.24 s and then stands for 0.24 s, always with two feet on the ground.
)
def predict_poses(pose, count):
predicted = [pose]
for _ in range(count):
x, y, yaw = predicted[-1]
foothold(name, period) answers where a foot should land at the end of its swing in a given period. The answer is
the spot that puts the foot at its nominal position under the base when the base passes over it in the middle of
the following stance, using the base poses predicted from the current velocity command (poses, below). A robot
that wants to go forward will therefore put its feet ahead of where they are, proportionally to the command, and one
that wants to turn will land its feet on an arc. For periods that are already over (pose_stage has moved past
them), the foothold is where the foot actually landed, kept in landed.
blend = 0.5 * (1 - np.cos(np.pi * phase))
return (1 - blend) * JUMP_LAND_HEIGHT + blend * base_height
def jump_foothold(name):
# Where the foot stands through the jump, which is where its last swing
# before the jump lands
period = jump_start // GAIT_PERIOD - 1
if name not in SWING_PAIRS[period % 2]:
period -= 1
return foothold(name, period)
def in_flight(k):
# Whether all feet are in the air at stage k
j = jump_stage(k)
flight = JUMP_CROUCH_STAGES + JUMP_PUSH_STAGES
return j is not None and flight <= j < flight + JUMP_FLIGHT_STAGES
foot_reference(name, k) is what the model needs: given a foot and a plan stage \(k\) (which may be anywhere in the
future), the position the foot should be at and whether it is swinging. A foot that is not in the swing pair for
that period stays at the foothold of its last swing. A swinging foot interpolates from its previous foothold to the
next with a cosine blend, which gives it zero velocity at both ends so that it lifts off and touches down without
sliding, and a raised-sine bump of step_height on top of that clears the ground in the middle of the swing.
The reference trajectory¶
# Base orientation relative to the walking direction
return rotation_z(yaw) @ rotation_y(-pitch) @ rotation_x(roll)
def stance_config(height, pitch, roll, yaw):
# Joint angles keeping each foot on its nominal spot on the ground, with
# the base at the given height and orientation
R = base_rotation(pitch, roll, yaw)
base = np.array([0.0, 0.0, height])
return np.hstack(
[
leg_ik(
R.T @ (foot_offsets[name] - base) - hip_offsets[name],
1 if foot_offsets[name][1] > 0 else -1,
)
predict_poses extends the base's pose \((x, y, \text{yaw})\) forward for count stages by integrating the current
velocity command, as if it were held constant. The command is in the base frame, so each step rotates it by the yaw.
These poses are used both as the base's reference, and, through foothold, to decide where the feet land. The
planner looks mpc_H + 3 * GAIT_PERIOD stages ahead rather than just mpc_H, since a foot swinging at the end of
the horizon needs the base pose from the stance after its swing to know where to land.
Costs and constraints¶
# Returns the foot's reference position at stage k and whether it swings
if jump_stage(k) is not None:
return jump_foothold(name), in_flight(k)
period, phase = divmod(k, GAIT_PERIOD)
phase /= GAIT_PERIOD
if name not in SWING_PAIRS[period % 2]:
return foothold(name, period - 1), False
# Ease in and out of the swing, so that the foot lifts off and touches
# down without sliding
blend = 0.5 * (1 - np.cos(np.pi * phase))
p = (1 - blend) * foothold(name, period - 2) + blend * foothold(
name, period
)
p[2] += step_height * np.sin(np.pi * phase) ** 2
return p, True
config_cost = ssqpy.model.costs.SquaredConfigurationErrorCost(
model, init_state[7:nq], C_WGT
)
v_cost = ssqpy.model.costs.SquaredJointVelocityCost(model, V_WGT)
u_cost = ssqpy.model.costs.SquaredControlCost(model, U_WGT)
f_cost = ssqpy.model.costs.SquaredForceCost(model, F_WGT)
base_costs = []
flight_costs = []
foot_funcs = {name: [] for name in foot_names}
for k in range(mpc_H):
model.addCost(k, config_cost)
model.addCost(k, v_cost)
model.addCost(k, u_cost)
model.addCost(k, f_cost)
base_costs.append(
ssqpy.model.costs.FrameSquaredPoseErrorCost(
model, "base_link", base_pos, R, T_WGT, R_WGT
)
)
model.addCost(k, base_costs[-1])
flight_costs.append(
ssqpy.model.costs.SquaredConfigurationErrorCost(
model,
stance_config(JUMP_TAKEOFF_HEIGHT, 0.0, 0.0, 0.0),
JUMP_CONFIG_WGT,
The shared costs (config_cost, v_cost, u_cost, f_cost) are single objects added to every stage, since they do
not change from stage to stage. The remaining ones are created per stage and kept in lists, so that
update_model can later retarget the one for stage idx independently:
-
base_costs[idx]is aFrameSquaredPoseErrorCostonbase_link, tracking the commanded base pose. -
foot_funcs[name][idx]holds, for each foot, aFrameSquaredTranslationErrorCostand the foot's contact constraint. The contact constraint has a friction coefficient of 0.6, a 150 N normal force limit, andno_slip=False, so that stance feet are not forced to a zero velocity. -
flight_costs[idx]is a configuration cost that is only used for the jump.
Every foot has both a contact constraint and a tracking cost at every stage; which of them is actually active
depends on whether the foot is swinging, and is decided at run time (below). This is the contact scheduling from
Go2 Jump, done with set_enabled on constraints and costs that
already exist, rather than by choosing which ones to add.
Reading the joystick¶
global jump_start, jump_button_held
# Forget the jump once it is over
if (
jump_start is not None
and current_stage >= jump_start + JUMP_PERIODS * GAIT_PERIOD
):
jump_start = None
# Schedule a jump on each new press of the jump button
held = JUMP_BUTTON in joystick.buttons
if (
held
and not jump_button_held
and jump_start is None
and np.max(np.abs(v_cmd)) < JUMP_MAX_VEL
):
jump_start = (current_stage // GAIT_PERIOD + 1) * GAIT_PERIOD
jump_button_held = held
target = np.array(
joystick.read(MAX_VEL_FRONT, MAX_VEL_BACK, MAX_VEL_SIDE, MAX_VEL_YAW)
)
if jump_start is not None:
target[:] = 0.0
max_change = np.array([MAX_ACC, MAX_ACC, MAX_YAW_ACC]) * mpc_dt
v_cmd[:] += np.clip(target - v_cmd, -max_change, max_change)
base_height = np.clip(
base_height - joystick.dpad_y * HEIGHT_RATE * mpc_dt,
MIN_BASE_HEIGHT,
MAX_BASE_HEIGHT,
)
step_height = np.clip(
step_height + joystick.dpad_x * STEP_HEIGHT_RATE * mpc_dt,
MIN_STEP_HEIGHT,
MAX_STEP_HEIGHT,
)
tilt = (PITCH_UP_BUTTON in joystick.buttons) - (
PITCH_DOWN_BUTTON in joystick.buttons
)
update_command is called once per stage and turns the gamepad state into the quantities the planner uses. The
stick command is first read, scaled to the maximum speeds MAX_VEL_FRONT, MAX_VEL_SIDE and so on, and then
v_cmd moves toward it by at most MAX_ACC * mpc_dt per stage, so the robot cannot be asked to accelerate faster
than it can follow. The remaining inputs are rate-limited in the same way, each one a value with a range and a rate:
the D-pad moves base_height and step_height, and pairs of buttons set base_pitch, base_roll and base_yaw to
a limit while held, and back to zero when released.
| Input | Effect |
|---|---|
| Sticks | Forward, sideways and yaw velocity command |
| D-pad up and down | Base height |
| D-pad left and right | Step height |
BTN_Z, BTN_TR |
Pitch the base up or down |
BTN_NORTH, BTN_EAST |
Roll the base to the left or right |
BTN_WEST, BTN_TL |
Turn the front to the left or right, with the feet staying in place |
BTN_C |
Jump in place, when the command is near rest |
The names are evdev codes, which map to different physical buttons on different gamepads.
Retargeting the horizon every stage¶
tilt * MAX_BASE_PITCH - base_pitch, -max_change, max_change
)
tilt = (ROLL_RIGHT_BUTTON in joystick.buttons) - (
ROLL_LEFT_BUTTON in joystick.buttons
)
max_change = ROLL_RATE * mpc_dt
base_roll += np.clip(
tilt * MAX_BASE_ROLL - base_roll, -max_change, max_change
)
turn = (YAW_LEFT_BUTTON in joystick.buttons) - (
YAW_RIGHT_BUTTON in joystick.buttons
)
max_change = YAW_RATE * mpc_dt
base_yaw += np.clip(
turn * MAX_BASE_YAW - base_yaw, -max_change, max_change
)
def update_model(current_stage):
global pose_stage, poses
# Feet finishing their swing keep the foothold they were aiming for
period = current_stage // GAIT_PERIOD
if current_stage > 0 and current_stage % GAIT_PERIOD == 0:
for name in SWING_PAIRS[(period - 1) % 2]:
landed[name] = foothold(name, period - 1)
# Look far enough ahead for the footholds of every swing in the horizon
poses = predict_poses(
poses[current_stage - pose_stage], mpc_H + 3 * GAIT_PERIOD
)
pose_stage = current_stage
config_cost.set_desired_configuration(
stance_config(base_height, base_pitch, base_roll, base_yaw)
)
for idx in range(mpc_H):
k = current_stage + idx
x, y, yaw = poses[idx]
j = jump_stage(k)
height = base_height if j is None else jump_base_height(j)
base_costs[idx].set_desired_translation(np.array([x, y, height]))
update_model is what makes the robot walk, and is called once per plan stage, after the joystick is read. It does
five things:
- Records landed feet. On the first stage of each new gait period, the pair that just finished swinging has
landed at its target foothold, and
landedis updated, so from now on the foot references use where the feet are, not where they were planned to land. - Rebuilds the base poses, starting from the pose the plan has reached, with the latest command.
- Sets the nominal configuration.
config_costis pointed atstance_config(...), the leg angles that keep every foot on its nominal spot with the base at the commanded height and orientation. This is how the D-pad and the tilt buttons work: they do not command anything directly, they change what "standing" means. - Sets the base cost of every stage to the pose for that stage and the tilt commands. The stage \(k\) of the model
corresponds to the absolute stage
current_stage + idxof the gait, which is whatfoot_referencetakes. - Sets each foot. Swinging feet have their contact constraint disabled and their cost tracking the swing path. Stance feet have the contact constraint enabled and a cost holding them at their foothold. As the comment in the code says, the constraint only bounds the foot's height, so the cost is also what stops stance feet sliding.
Note that the schedule is entirely periodic: at every stage the same function is evaluated at a different \(k\), and
the window slides along, with a foot's contact turning off as its swing enters the horizon and on again as it leaves
it. This is why no part of the model ever needs to be rebuilt after model.finalize.
The jump¶
return predicted
def foothold(name, period):
# Where the foot lands at the end of its swing in the given period, placed
# so that the base passes over the foot's nominal position mid-stance
if (period + 1) * GAIT_PERIOD <= pose_stage:
return landed[name]
x, y, yaw = poses[int((period + 1.5) * GAIT_PERIOD) - pose_stage]
return np.array([x, y, 0.0]) + rotation_z(yaw) @ foot_offsets[name]
def jump_stage(k):
# Stage k's index within the jump, or None outside of it
if jump_start is None:
return None
j = k - jump_start
return j if 0 <= j < JUMP_PERIODS * GAIT_PERIOD else None
def jump_base_height(j):
# Base reference height at stage j of the jump
push = JUMP_CROUCH_STAGES
flight = push + JUMP_PUSH_STAGES
Pressing the jump button when the command is near rest schedules a jump in place, starting at the next gait period.
It reuses the machinery above rather than adding a new one. During the jump (jump_stage(k) is not None):
-
The base height reference follows
jump_base_height: a crouch, a push-off with constant acceleration up to the takeoff speed, a ballistic flight, a landing that decelerates, and a recovery back to the normal height. The takeoff speedJUMP_SPEEDis set from the flight time, so that the base lands at the height it left. The velocity command is held at zero throughout. -
All four feet stand at the same footholds, those of their last swing before the jump (
jump_foothold). -
In flight (
in_flight(k)), every foot is marked as swinging, so all contact constraints and foot costs are disabled. The base's translation weight is zeroed too, since its path can no longer be steered once the feet are off the ground, and aflight_costsconfiguration cost is enabled instead to hold the legs in their takeoff pose.
This is the same stance, flight, stance structure as Go2 Jump, but built from the pieces that already exist for walking, and decided at every stage instead of once.
Solver parameters¶
for name in foot_names:
foot_pos, swinging = foot_reference(name, k)
foot_cost, foot_constraint = foot_funcs[name][idx]
# Swinging feet track their swing trajectory, while stance feet
# stay in contact. The contact constraint only bounds the foot's
# height, so the tracking cost also keeps stance feet from sliding
foot_constraint.set_enabled(not swinging)
foot_cost.set_enabled(not flight)
foot_cost.set_desired_translation(foot_pos)
update_model(0)
model.addCost(mpc_H, v_cost)
model.finalize(init_states=[init_state])
HPIPM is the QP solver, with tolerances of \(10^{-4}\). The MPC wrapper is given just two
SQP iterations, since the problem is warm started from the previous stage's solution (see
Warm-starting). The admm_params are set in the script but unused, as long as the
QP solver type is HPIPM.
The simulation loop¶
# Start the viewer maximized
glfw.init()
glfw.window_hint(glfw.MAXIMIZED, glfw.TRUE)
with mujoco.viewer.launch_passive(mjmodel, mjdata) as viewer:
while viewer.is_running():
display = step % DISPLAY_EVERY == 0
"""
if display:
print(f"\rSimulation time: {mjdata.time:.2f}s")
if step > 0:
loop_rate = DISPLAY_EVERY / (time() - display_time)
print(
f"\rLoop frequency: {loop_rate:.2f}Hz\t"
f"MPC frequency: {(loop_rate / STAGE_EVERY):.2f}Hz"
)
display_time = time()
"""
state = mujoco_to_ssqp_state(np.hstack((mjdata.qpos, mjdata.qvel)))
state = clip_joints_to_limits(state, mjmodel)
# Solve once per plan stage and hold the stage's torques in between,
# as the planner assumes
if step % STAGE_EVERY == 0:
try:
x, u = mpc.step(state)
xk = ssqp_to_mujoco_state(x.stage(1))
if dynamics_mode == ssqpy.model.SolverMode.ForwardDynamics:
uk = ssqp_to_mujoco_control(u.stage(0))
elif (
dynamics_mode
== ssqpy.model.SolverMode.ReducedInverseDynamics
):
uk = ssqp_to_mujoco_control(
np.hstack(
(
np.zeros(6),
model.getActuatedTorqueVector(
state, u.stage(0)
Every simulation step, the state is read from MuJoCo and converted (quaternion order, base velocity frame and joint
ordering, by utils_go2_mujoco.py) and clipped to the joint limits that the model enforces, because MuJoCo's joint
limits are soft. Then:
-
Every
STAGE_EVERYsteps,mpc.step(state)solves. From the result it keeps the next stage's statexk, as a joint target, and the torques that the first stage's accelerations and forces need at the current measured state, frommodel.getActuatedTorqueVector. In forward dynamics mode, the control is the torque already. If the solver fails with a regularization error, the loop exits. -
Every step, the command applied is those torques, held until the next solve, plus a small joint PD term (
Kp= 0.01) toward the next stage's joint position and velocity, which corrects drift between solves. The[6:18]slice drops the base entries, which are not actuated, andssqp_to_mujoco_ctrl_kinreorders the legs for MuJoCo's actuators. -
Every
STAGE_EVERYsteps, after the command is set,stageadvances andupdate_commandandupdate_modelrun, so the next solve sees the retargeted horizon. They run beforemj_step, in the same iteration.
The torques are therefore computed once per solve, from the state at that moment, and held for the 12 simulation steps of the stage. This matches the planner's assumption that the controls are constant over a stage.
Running it¶
From examples/mujoco_mpc/go2/, with a gamepad connected, the environment set up per
Installation and MuJoCo's own virtual environment
per Examples:
The MuJoCo viewer opens maximized and the robot stands. Use the sticks to walk, and the other controls to change the
posture. The script does not wait to match real time, so the simulation runs as fast as the solves allow, and
printing the frequency is available in the commented block at the top of the loop. The inputs package is needed to
read the gamepad.
Things to try¶
-
Change
GAIT_PERIODto make the trot faster or slower, and see how the speeds that work change with it. -
Raise
STEP_HEIGHT, or the base height, and check how the robot handles it. -
Lower
mpc_H, and notice that the footholds near the end of the horizon, and so the response to a command change, get less accurate. -
Set
Kpto zero to see how much the small PD term between solves contributes. -
Change
dynamics_modetoForwardDynamicsand retuneU_WGTandF_WGT: the weights are different for the two modes for a reason. -
Press the jump button while walking: it is ignored unless the velocity command is below
JUMP_MAX_VEL.