|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +from pathlib import Path |
| 4 | + |
| 5 | +import mujoco |
| 6 | +import numpy as np |
| 7 | +from PIL import Image, ImageDraw |
| 8 | + |
| 9 | + |
| 10 | +ROOT = Path(__file__).resolve().parents[5] |
| 11 | +FIG_DIR = ROOT / "docs/practices/quadruped/cs123/figs" |
| 12 | +FIG_DIR.mkdir(parents=True, exist_ok=True) |
| 13 | + |
| 14 | +LEGS = ("FL", "FR", "RL", "RR") |
| 15 | +LEG_Y = {"FL": 0.07, "FR": -0.07, "RL": 0.07, "RR": -0.07} |
| 16 | +LEG_X = {"FL": 0.105, "FR": 0.105, "RL": -0.105, "RR": -0.105} |
| 17 | +PHASE_OFFSETS = {"FL": 0.0, "FR": 0.5, "RL": 0.5, "RR": 0.0} |
| 18 | + |
| 19 | +THIGH = 0.105 |
| 20 | +CALF = 0.105 |
| 21 | +T_CYCLE = 0.4 |
| 22 | +STEP_HEIGHT = 0.04 |
| 23 | +STAND_HEIGHT = 0.17 |
| 24 | + |
| 25 | + |
| 26 | +def build_xml() -> str: |
| 27 | + legs = [] |
| 28 | + for leg in LEGS: |
| 29 | + y = LEG_Y[leg] |
| 30 | + x = LEG_X[leg] |
| 31 | + side = 1 if y > 0 else -1 |
| 32 | + rgba = "0.10 0.55 0.20 1" if leg in ("FL", "RR") else "0.10 0.32 0.75 1" |
| 33 | + legs.append( |
| 34 | + f""" |
| 35 | + <body name="{leg}_hip" pos="{x:.3f} {y:.3f} 0"> |
| 36 | + <joint name="{leg}_hip" type="hinge" axis="1 0 0" range="-0.8 0.8" damping="1.0" armature="0.01"/> |
| 37 | + <geom type="sphere" size="0.017" rgba="{rgba}"/> |
| 38 | + <body name="{leg}_thigh" pos="0 {0.018 * side:.3f} 0"> |
| 39 | + <joint name="{leg}_thigh" type="hinge" axis="0 1 0" range="-1.8 1.8" damping="1.0" armature="0.01"/> |
| 40 | + <geom type="capsule" fromto="0 0 0 0 0 {-THIGH:.3f}" size="0.011" rgba="0.55 0.55 0.55 1"/> |
| 41 | + <body name="{leg}_calf" pos="0 0 {-THIGH:.3f}"> |
| 42 | + <joint name="{leg}_calf" type="hinge" axis="0 1 0" range="-2.4 0.2" damping="1.0" armature="0.01"/> |
| 43 | + <geom type="capsule" fromto="0 0 0 0 0 {-CALF:.3f}" size="0.009" rgba="{rgba}"/> |
| 44 | + <body name="{leg}_foot" pos="0 0 {-CALF:.3f}"> |
| 45 | + <geom name="{leg}_foot_geom" type="sphere" size="0.018" friction="1.6 0.02 0.002" rgba="0.07 0.07 0.07 1"/> |
| 46 | + </body> |
| 47 | + </body> |
| 48 | + </body> |
| 49 | + </body> |
| 50 | + """ |
| 51 | + ) |
| 52 | + |
| 53 | + return f""" |
| 54 | + <mujoco model="cs123_gait_demo"> |
| 55 | + <compiler angle="radian"/> |
| 56 | + <option timestep="0.002" gravity="0 0 -9.81"/> |
| 57 | + <visual> |
| 58 | + <global offwidth="720" offheight="540"/> |
| 59 | + <headlight diffuse="0.7 0.7 0.7" ambient="0.35 0.35 0.35"/> |
| 60 | + <quality shadowsize="2048"/> |
| 61 | + </visual> |
| 62 | + <asset> |
| 63 | + <texture name="grid" type="2d" builtin="checker" rgb1="0.92 0.94 0.92" rgb2="0.78 0.83 0.78" width="512" height="512"/> |
| 64 | + <material name="ground" texture="grid" texrepeat="4 4" reflectance="0.12"/> |
| 65 | + </asset> |
| 66 | + <worldbody> |
| 67 | + <light pos="0 -2 3" dir="0 1 -1" diffuse="0.8 0.8 0.8"/> |
| 68 | + <geom name="floor" type="plane" size="8 8 0.1" material="ground"/> |
| 69 | + <body name="torso" pos="0 0 {STAND_HEIGHT:.3f}"> |
| 70 | + <freejoint/> |
| 71 | + <geom name="torso_geom" type="box" size="0.135 0.052 0.035" rgba="0.72 0.72 0.72 1" mass="2.0"/> |
| 72 | + {''.join(legs)} |
| 73 | + </body> |
| 74 | + </worldbody> |
| 75 | + <actuator> |
| 76 | + {''.join(f'<motor joint="{leg}_{joint}" gear="1"/>' for leg in LEGS for joint in ("hip", "thigh", "calf"))} |
| 77 | + </actuator> |
| 78 | + </mujoco> |
| 79 | + """ |
| 80 | + |
| 81 | + |
| 82 | +def leg_phase(t: float, leg: str, t_cycle: float = T_CYCLE, duty: float = 0.5) -> tuple[bool, float]: |
| 83 | + t_global = (t / t_cycle) % 1.0 |
| 84 | + t_local = (t_global + PHASE_OFFSETS[leg]) % 1.0 |
| 85 | + if t_local < duty: |
| 86 | + return True, t_local / duty |
| 87 | + return False, (t_local - duty) / (1.0 - duty) |
| 88 | + |
| 89 | + |
| 90 | +def foot_trajectory(s: float, in_stance: bool, step_length: float) -> np.ndarray: |
| 91 | + if in_stance: |
| 92 | + x = step_length * (0.5 - s) |
| 93 | + z = -STAND_HEIGHT |
| 94 | + else: |
| 95 | + x = step_length * (s - 0.5) |
| 96 | + z = -STAND_HEIGHT + STEP_HEIGHT * np.sin(np.pi * s) |
| 97 | + return np.array([x, z]) |
| 98 | + |
| 99 | + |
| 100 | +def ik_leg_2d(x: float, z: float) -> tuple[float, float]: |
| 101 | + down = -z |
| 102 | + reach = np.hypot(x, down) |
| 103 | + reach = np.clip(reach, 0.045, THIGH + CALF - 0.004) |
| 104 | + cos_knee = (reach * reach - THIGH * THIGH - CALF * CALF) / (2 * THIGH * CALF) |
| 105 | + cos_knee = np.clip(cos_knee, -0.98, 0.98) |
| 106 | + knee = -np.arccos(cos_knee) |
| 107 | + hip = np.arctan2(x, down) - np.arctan2(CALF * np.sin(knee), THIGH + CALF * np.cos(knee)) |
| 108 | + return hip, knee |
| 109 | + |
| 110 | + |
| 111 | +def gait_step(t: float, step_length: float) -> np.ndarray: |
| 112 | + target = np.zeros(12) |
| 113 | + for i, leg in enumerate(LEGS): |
| 114 | + in_stance, s = leg_phase(t, leg) |
| 115 | + foot_xz = foot_trajectory(s, in_stance, step_length) |
| 116 | + thigh, calf = ik_leg_2d(float(foot_xz[0]), float(foot_xz[1])) |
| 117 | + target[3 * i : 3 * i + 3] = [0.0, thigh, calf] |
| 118 | + return target |
| 119 | + |
| 120 | + |
| 121 | +def add_label(frame: np.ndarray, text: str) -> Image.Image: |
| 122 | + image = Image.fromarray(frame) |
| 123 | + draw = ImageDraw.Draw(image, "RGBA") |
| 124 | + draw.rounded_rectangle((18, 18, 330, 68), radius=10, fill=(255, 255, 255, 210)) |
| 125 | + draw.text((34, 34), text, fill=(20, 30, 40, 255)) |
| 126 | + return image |
| 127 | + |
| 128 | + |
| 129 | +def render_experiment(name: str, output: Path, step_length: float, base_speed: float) -> None: |
| 130 | + model = mujoco.MjModel.from_xml_string(build_xml()) |
| 131 | + data = mujoco.MjData(model) |
| 132 | + renderer = mujoco.Renderer(model, height=540, width=720) |
| 133 | + |
| 134 | + camera = mujoco.MjvCamera() |
| 135 | + camera.lookat[:] = [0.18, 0.0, 0.09] |
| 136 | + camera.distance = 0.78 |
| 137 | + camera.azimuth = 135 |
| 138 | + camera.elevation = -18 |
| 139 | + |
| 140 | + q0 = gait_step(0.0, step_length) |
| 141 | + data.qpos[:7] = [0.0, 0.0, STAND_HEIGHT, 1.0, 0.0, 0.0, 0.0] |
| 142 | + data.qpos[7:] = q0 |
| 143 | + mujoco.mj_forward(model, data) |
| 144 | + |
| 145 | + duration = 3.2 |
| 146 | + fps = 24 |
| 147 | + frames: list[Image.Image] = [] |
| 148 | + |
| 149 | + for frame_index in range(int(duration * fps)): |
| 150 | + t = frame_index / fps |
| 151 | + q_des = gait_step(t, step_length) |
| 152 | + # Keep the torso path prescribed so the GIF focuses on gait timing rather than |
| 153 | + # on model-specific balance tuning. |
| 154 | + data.qpos[:7] = [base_speed * t, 0.0, STAND_HEIGHT, 1.0, 0.0, 0.0, 0.0] |
| 155 | + data.qpos[7:] = q_des |
| 156 | + data.qvel[:] = 0.0 |
| 157 | + mujoco.mj_forward(model, data) |
| 158 | + |
| 159 | + camera.lookat[0] = base_speed * t + 0.05 |
| 160 | + renderer.update_scene(data, camera=camera) |
| 161 | + frames.append(add_label(renderer.render(), name)) |
| 162 | + |
| 163 | + frames[0].save( |
| 164 | + output, |
| 165 | + save_all=True, |
| 166 | + append_images=frames[1:], |
| 167 | + duration=1000 // fps, |
| 168 | + loop=0, |
| 169 | + optimize=True, |
| 170 | + ) |
| 171 | + print(f"saved {output} ({len(frames)} frames)") |
| 172 | + |
| 173 | + |
| 174 | +def main() -> None: |
| 175 | + render_experiment( |
| 176 | + name="Experiment 1: in-place trot", |
| 177 | + output=FIG_DIR / "lab5_inplace_trot.gif", |
| 178 | + step_length=0.0, |
| 179 | + base_speed=0.0, |
| 180 | + ) |
| 181 | + render_experiment( |
| 182 | + name="Experiment 2: forward trot", |
| 183 | + output=FIG_DIR / "lab5_forward_trot.gif", |
| 184 | + step_length=0.10, |
| 185 | + base_speed=0.22, |
| 186 | + ) |
| 187 | + |
| 188 | + |
| 189 | +if __name__ == "__main__": |
| 190 | + main() |
0 commit comments