diff --git a/biped_walking_controller/foot.py b/biped_walking_controller/foot.py index afe489a..3d9d702 100644 --- a/biped_walking_controller/foot.py +++ b/biped_walking_controller/foot.py @@ -195,7 +195,28 @@ def compute_steps_sequence( for i in range(1, n_steps + 1): sign = -1.0 if i % 2 == 0 else 1.0 steps_ids.append(Foot.RIGHT if i % 2 == 0 else Foot.LEFT) - steps_pose[i] = np.array([i * l_stride, sign * math.fabs(lf_initial_pose[1]), 0.0]) + steps_pose[i] = np.array( + [i * l_stride, sign * math.fabs(lf_initial_pose[1]), rf_initial_pose[2]] + ) + + # Add a last step to have both feet at the same level + steps_pose[-1] = steps_pose[-2] + steps_pose[-1][1] = steps_pose[-1][1] * -1.0 + + return steps_pose, steps_ids + + +def compute_onplace_steps_sequence( + rf_initial_pose: np.ndarray, + lf_initial_pose: np.ndarray, + n_steps: int, +): + steps_pose = np.zeros((n_steps + 2, 3)) + steps_pose[0] = rf_initial_pose + steps_ids = [Foot.RIGHT] + for i in range(1, n_steps + 1): + steps_ids.append(Foot.RIGHT if i % 2 == 0 else Foot.LEFT) + steps_pose[i] = rf_initial_pose if i % 2 == 0 else lf_initial_pose # Add a last step to have both feet at the same level steps_pose[-1] = steps_pose[-2] diff --git a/biped_walking_controller/plot.py b/biped_walking_controller/plot.py index 80db9e7..a136c23 100644 --- a/biped_walking_controller/plot.py +++ b/biped_walking_controller/plot.py @@ -2,7 +2,6 @@ compute_double_support_polygon, compute_single_support_polygon, ) -from biped_walking_controller.state_machine import WalkingState def plot_steps(axes, steps_pose, step_shape): @@ -145,6 +144,7 @@ def plot_feet_and_com( ax.set_ylabel(coord_labels[coord_idx.index(j)]) ax.grid(True) + axes[1].set_ylim((-0.02, 0.06)) axes[-1].set_xlabel("t [s]") # One combined legend outside @@ -187,36 +187,20 @@ def traj2d(arr): # drop last 2 samples as in your code ax2.set_title(f"{title_prefix} — plan view (x–y)") -def plot_contact_forces_and_state( +def plot_contact_forces( t: float, force_rf: float, force_lf: float, - states: WalkingState, - title: str = "Contact force Fx", ): t = np.asarray(t).ravel() force_rf = np.asarray(force_rf).ravel() force_lf = np.asarray(force_lf).ravel() assert t.size == force_rf.size == force_lf.size - fig, ax = plt.subplots(2, figsize=(12, 8), layout="constrained", sharex=True) - ax[0].plot(t, force_rf, label="right foot", marker="*") - ax[0].plot(t, force_lf, label="left foot", marker="*") - ax[0].set_ylabel("Normal force [N]") - ax[0].set_title("Contact Force Fx") - ax[0].grid(True) - ax[0].legend() - - ticks = [] - labels = [] - for data in WalkingState: - ticks.append(data.value) - labels.append(data.name) - - ax[1].plot(t, states, label="state", marker="*") - ax[1].set_xlabel("t [s]") - ax[1].set_ylabel("State [-]") - ax[1].set_title("State") - ax[1].set_yticks(ticks, labels) - ax[1].grid(True) - ax[1].legend() + fig, ax = plt.subplots(figsize=(12, 8), layout="constrained", sharex=True) + ax.plot(t, force_rf, label="right foot", marker="*") + ax.plot(t, force_lf, label="left foot", marker="*") + ax.set_ylabel("Normal force [N]") + ax.set_title("Contact Force Fx") + ax.grid(True) + ax.legend() diff --git a/biped_walking_controller/preview_control.py b/biped_walking_controller/preview_control.py index b0c10b1..d81d4e6 100644 --- a/biped_walking_controller/preview_control.py +++ b/biped_walking_controller/preview_control.py @@ -1,12 +1,11 @@ -import abc import typing -from abc import abstractmethod from dataclasses import dataclass -from enum import Enum import numpy as np from scipy.linalg import solve_discrete_are +from biped_walking_controller.state_machine import WalkingState, Foot + def linear_interpolation(t: np.ndarray, pos_begin: np.ndarray, pos_end: np.ndarray) -> np.ndarray: return (1 - t)[:, None] * pos_begin + t[:, None] * pos_end @@ -15,6 +14,13 @@ def linear_interpolation(t: np.ndarray, pos_begin: np.ndarray, pos_end: np.ndarr def cubic_spline_interpolation( t: np.ndarray, pos_begin: np.ndarray, pos_end: np.ndarray ) -> np.ndarray: + if isinstance(t, float): + return ( + pos_begin + + 3.0 * (pos_end - pos_begin) * np.square(t) + - 2.0 * (pos_end - pos_begin) * np.pow(t, 3) + ) + return ( pos_begin + 3.0 * (pos_end - pos_begin) * np.square(t)[:, None] @@ -255,3 +261,243 @@ def update_control(ctrl_mat: PreviewControllerMatrices, current_zmp, zmp_ref, x, y_next[1:] = ctrl_mat.A @ y[1:] + ctrl_mat.B.ravel() * u[1] return u, x_next, y_next + + +def build_zmp_horizon( + com_initial_target, + t_horizon: float, + t_state: float, + state: WalkingState, + delta_t: float, + current_step_idx: int, + steps_sequence: np.ndarray, + steps_foot: typing.List[Foot], + ss_t: float, + ds_t: float, + t_init: float, + t_end: float, + interp_fn=cubic_spline_interpolation, +) -> tuple[np.ndarray, np.ndarray]: + """ + Build a ZMP reference over a preview horizon based on the current walking state. + + Parameters + ---------- + t_horizon : float + Length of the preview horizon [s]. + t_state : float + Time elapsed in the current state [s]. + state : WalkingState + Current walking state (INIT, DS, SS_LEFT, SS_RIGHT, END). + delta_t : float + Sampling period of the preview horizon [s]. + current_step_idx : int + Index of the current step in `steps_sequence`. Interpreted as: + - In SS: index of the current support foot. + - In DS: index of the *target* foot (previous is current_step_idx - 1). + steps_sequence : np.ndarray + Sequence of footsteps. Shape (N, 2) or (N, >=2). + Only the (x, y) components are used as ZMP targets. + ss_t : float + Duration of a single-support phase [s]. + ds_t : float + Duration of a double-support phase [s]. + t_init : float + Duration of the INIT phase [s]. + t_end : float + Duration of the END phase [s]. + interp_fn : callable + Interpolation function used in double support. + Expected signature: interp_fn(alpha, p0, p1) + where alpha in [0, 1], p0, p1 are 2D numpy arrays. + + Returns + ------- + t_samples : np.ndarray, shape (N,) + Relative time samples over the horizon, starting at 0. + zmp_horizon : np.ndarray, shape (N, 2) + ZMP reference (x, y) at each time sample. + + Notes + ----- + - This function does NOT enforce continuity with the *previous* ZMP reference. + At each call, the horizon is built from scratch from the current state. + - State-transition logic is a simple cyclic model: + INIT -> DS -> SS -> DS -> SS -> ... -> END + and the step index increments when leaving an SS state. + """ + # Sanity checks and normalization + steps = np.asarray(steps_sequence, dtype=float) + if steps.ndim != 2 or steps.shape[0] == 0: + raise ValueError("steps_sequence must be a non-empty (N, D) array") + + # Use only x, y + steps_xy = steps[:, :2] + n_steps = steps_xy.shape[0] + + step_idx = int(np.clip(current_step_idx, 0, n_steps - 1)) + + def state_duration(s: WalkingState) -> float: + if s == WalkingState.INIT: + return t_init + if s == WalkingState.DS: + return ds_t + if s in (WalkingState.SS_LEFT, WalkingState.SS_RIGHT): + return ss_t + if s == WalkingState.END: + return t_end + raise ValueError(f"Unknown state: {s}") + + def next_state_and_step( + s: WalkingState, idx: int, steps_foot: typing.List[Foot] + ) -> tuple[WalkingState, int]: + """Simple progression model over footsteps. + + INIT -> DS(step 0) + DS(k) -> SS(k) + SS(k) -> DS(k+1) (clamped to last step) + END stays END + """ + if s == WalkingState.DS or s == WalkingState.INIT: + # Land on target step + if idx < n_steps - 1: + return ( + WalkingState.SS_LEFT if steps_foot[idx] is Foot.LEFT else WalkingState.SS_RIGHT + ), idx + else: + return WalkingState.END, idx + + if s in (WalkingState.SS_LEFT, WalkingState.SS_RIGHT): + # Move to next DS between current and next step + if idx < n_steps - 1: + return WalkingState.DS, idx + 1 + + if s == WalkingState.END: + return WalkingState.END, idx + + raise ValueError(f"Unknown state: {s}") + + def zmp_for_state(s: WalkingState, idx: int, t_in_state: float) -> np.ndarray: + """Return ZMP position (x, y) for a given state and time in state.""" + # Clamp index + idx = int(np.clip(idx, 0, n_steps - 1)) + + if s == WalkingState.INIT: + # Keep ZMP on the first support + p0 = com_initial_target + p1 = steps_xy[0] + + alpha = np.clip(t_in_state / t_init, 0.0, 1.0) + return interp_fn(alpha, p0, p1) + + if s in (WalkingState.SS_LEFT, WalkingState.SS_RIGHT): + # ZMP fixed at current support foot + return steps_xy[idx] + + if s == WalkingState.DS: + # ZMP moves between previous and current step + if idx <= 0: + p0 = com_initial_target + p1 = steps_xy[0] + else: + p0 = steps_xy[idx - 1] + p1 = steps_xy[idx] + + alpha = np.clip(t_in_state / max(ds_t, 1e-6), 0.0, 1.0) + return interp_fn(alpha, p0, p1) + + if s == WalkingState.END: + # Keep ZMP on last support + return (steps_xy[-2] + steps_xy[-1]) / 2.0 + + raise ValueError(f"Unknown state: {s}") + + # Allocate horizon + n_samples = int(np.floor(t_horizon / delta_t)) + t_samples = np.arange(n_samples, dtype=float) * delta_t + zmp_horizon = np.zeros((n_samples, 2), dtype=float) + + # Simulated state over the horizon + sim_state = state + time_in_state = t_state + sim_step_idx = step_idx + + for k in range(n_samples): + # Advance state machine if needed (except END which just saturates) + while sim_state != WalkingState.END and time_in_state >= state_duration(sim_state): + time_in_state -= state_duration(sim_state) + sim_state, sim_step_idx = next_state_and_step(sim_state, sim_step_idx, steps_foot) + + # Compute ZMP for current simulated state + zmp_horizon[k] = zmp_for_state(sim_state, sim_step_idx, time_in_state) + + # Advance local time + time_in_state += delta_t + + return t_samples, zmp_horizon + + +class CentroidalPlanner: + def __init__( + self, + dt: float, + com_initial_target: np.ndarray, + params: PreviewControllerParams, + ): + self.params = params + self.dt = dt + self.ctrler_mat = compute_preview_control_matrices(params, dt) + self.x = np.array([0.0, com_initial_target[0], 0.0, 0.0], dtype=float) + self.y = np.array([0.0, com_initial_target[1], 0.0, 0.0], dtype=float) + self.steps_sequence = None + self.steps_foot = None + self.step_idx = 0 + self.zmp_ref_horizon = None + + def set_steps_sequence(self, steps_sequence: np.ndarray, steps_foot: typing.List[Foot]): + # We assign the new sequence and reset the step counter + self.steps_sequence = steps_sequence + self.steps_foot = steps_foot + self.step_idx = 0 + + def update( + self, + com_initial_target, + state: WalkingState, + step_idx: int, + t_state: float, + t_init: float, + t_end: float, + t_ss: float, + t_ds: float, + ): + # Update control + _, self.zmp_ref_horizon = build_zmp_horizon( + com_initial_target, + t_horizon=(self.params.n_preview_steps - 1) * self.dt, + t_state=t_state, + state=state, + delta_t=self.dt, + current_step_idx=step_idx, + steps_sequence=self.steps_sequence, + steps_foot=self.steps_foot, + ss_t=t_ss, + ds_t=t_ds, + t_init=t_init, + t_end=t_end, + interp_fn=cubic_spline_interpolation, + ) + + _, self.x, self.y = update_control( + self.ctrler_mat, + self.zmp_ref_horizon[0], + self.zmp_ref_horizon, + self.x.copy(), + self.y.copy(), + ) + + def get_com_pos(self) -> typing.Tuple[float, float]: + return self.x[1], self.y[1] + + def get_ref_horizon(self): + return self.zmp_ref_horizon diff --git a/biped_walking_controller/simulation.py b/biped_walking_controller/simulation.py index d2e708c..831c96c 100644 --- a/biped_walking_controller/simulation.py +++ b/biped_walking_controller/simulation.py @@ -1,3 +1,5 @@ +import os +import subprocess import typing from pathlib import Path from typing import Tuple, Dict @@ -340,6 +342,50 @@ def _compute_force( return total_force, [mean_x, mean_y, mean_z] +def _process_video(filename: str, simulation_duration: float, start_time: float = 0.0): + cmd = [ + "ffprobe", + "-i", + filename, + "-show_entries", + "format=duration", + "-v", + "quiet", + "-of", + "csv=p=0", + ] + + out = subprocess.check_output(cmd).decode().strip() + video_duration = float(out) + + ratio = simulation_duration / video_duration + cmd = [ + "ffmpeg", + "-y", + "-ss", + str(start_time / ratio), + "-i", + filename, + "-filter:v", + f"setpts={ratio}*PTS", + "-r", + "60", + "-c:v", + "libx264", + "-crf", + "18", + "-preset", + "slow", + "output.mp4", + ] + + subprocess.run(cmd, check=True) + + # Remove the original file and rename the output to the original file + os.remove(filename) + os.rename("output.mp4", filename) + + class Simulator: """ Thin PyBullet wrapper for loading a robot URDF and driving it from @@ -348,7 +394,14 @@ class Simulator: The robot is loaded at the origin and a flat horizontal plane is added at z=0. """ - def __init__(self, dt, path_to_robot_urdf: Path, model, launch_gui=True): + def __init__( + self, + dt, + path_to_robot_urdf: Path, + model, + launch_gui=True, + n_solver_iter=200, + ): """ Initialize PyBullet, load ground and TALOS, and set physics. @@ -375,17 +428,22 @@ def __init__(self, dt, path_to_robot_urdf: Path, model, launch_gui=True): pb.GUI if launch_gui else pb.DIRECT, options="--window_title=PyBullet --width=1920 --height=1080", ) + + pb.configureDebugVisualizer(pb.COV_ENABLE_GUI, 0) + pb.configureDebugVisualizer(pb.COV_ENABLE_RGB_BUFFER_PREVIEW, 0) + pb.configureDebugVisualizer(pb.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0) + pb.configureDebugVisualizer(pb.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0) pb.setAdditionalSearchPath(pybullet_data.getDataPath()) pb.setGravity(0, 0, -9.81) pb.setTimeStep(dt) pb.setRealTimeSimulation(0) pb.setPhysicsEngineParameter( fixedTimeStep=dt, - numSolverIterations=100, + numSolverIterations=n_solver_iter, numSubSteps=1, useSplitImpulse=1, splitImpulsePenetrationThreshold=0.01, - contactSlop=0.001, + contactSlop=0.005, erp=0.2, contactERP=0.2, frictionERP=0.05, @@ -417,6 +475,8 @@ def __init__(self, dt, path_to_robot_urdf: Path, model, launch_gui=True): self._displayed_lines = None self._displayed_points = None + self.log_id = None + self.filename = None def step(self): """ @@ -497,7 +557,7 @@ def update_camera_to_follow_pos(self, x: float, y: float, z: float): """ pb.resetDebugVisualizerCamera( cameraDistance=4.0, - cameraYaw=50, + cameraYaw=120, cameraPitch=-40, cameraTargetPosition=[x, y, z], ) @@ -718,3 +778,12 @@ def get_zmp_pose(self): px = -M[1] / F[2] py = M[0] / F[2] return np.array([px, py, 0.0]) + + def start_video_record(self, filename: str = "recording.mp4"): + self.filename = filename + self.log_id = pb.startStateLogging(pb.STATE_LOGGING_VIDEO_MP4, self.filename) + + def stop_video_record(self, duration: float, t_start: float = 0.0): + pb.stopStateLogging(self.log_id) + + _process_video(self.filename, duration, t_start) diff --git a/examples/example_4_physics_simulation.py b/examples/example_4_physics_simulation.py index 3b2435b..fc04d31 100644 --- a/examples/example_4_physics_simulation.py +++ b/examples/example_4_physics_simulation.py @@ -1,5 +1,7 @@ import math import argparse +import typing +from dataclasses import dataclass from pathlib import Path import numpy as np @@ -14,7 +16,7 @@ ) from biped_walking_controller.inverse_kinematic import InvKinSolverParams, solve_inverse_kinematics -from biped_walking_controller.plot import plot_feet_and_com, plot_contact_forces_and_state +from biped_walking_controller.plot import plot_feet_and_com, plot_contact_forces from biped_walking_controller.preview_control import ( PreviewControllerParams, @@ -33,18 +35,9 @@ ) -def main(): - p = argparse.ArgumentParser() - p.add_argument("--path-talos-data", type=Path, help="Path to talos_data root") - p.add_argument("--plot-results", action="store_true") - p.add_argument("--launch-gui", action="store_true") - args = p.parse_args() - - np.set_printoptions(suppress=True, precision=3) - +@dataclass +class GeneralParams: dt = 1.0 / 240.0 - - # ZMP reference parameters t_ss = 0.8 # Single support phase time window t_ds = 0.3 # Double support phase time window t_init = 2.0 # Initialization phase (transition from still position to first step) @@ -52,18 +45,55 @@ def main(): n_steps = 15 # Number of steps executed by the robot l_stride = 0.1 # Length of the stride max_height_foot = 0.01 # Maximal height of the swing foot + t_preview = 1.6 + n_solver_iter = 200 + + +def get_standard_params() -> typing.Tuple[GeneralParams, PreviewControllerParams]: + general_params = GeneralParams() - # Preview controller parameters - t_preview = 1.6 # Time horizon used for the preview controller ctrler_params = PreviewControllerParams( zc=0.89, g=9.81, Qe=1.0, Qx=np.zeros((3, 3)), R=1e-6, - n_preview_steps=int(round(t_preview / dt)), + n_preview_steps=int(round(general_params.t_preview / general_params.dt)), ) - ctrler_mat = compute_preview_control_matrices(ctrler_params, dt) + + return general_params, ctrler_params + + +def get_accurate_sim_params() -> typing.Tuple[GeneralParams, PreviewControllerParams]: + general_params, ctrler_params = get_standard_params() + + # Specific params + general_params.dt = 1.0 / 1000.0 + general_params.t_ss = 0.6 + general_params.t_ds = 0.1 + general_params.n_steps = 15 + general_params.l_stride = 0.15 + general_params.max_height_foot = 0.02 + general_params.n_solver_iter = 1500 + + return general_params, ctrler_params + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--path-talos-data", type=Path, help="Path to talos_data root") + p.add_argument("--plot-results", action="store_true") + p.add_argument("--launch-gui", action="store_true") + p.add_argument("--record-video", action="store_true") + p.add_argument("--accurate-sim", action="store_true") + args = p.parse_args() + + np.set_printoptions(suppress=True, precision=3) + + if args.accurate_sim: + scen_params, ctrler_params = get_accurate_sim_params() + else: + scen_params, ctrler_params = get_standard_params() # Initialize Talos pinocchio model talos = Talos(path_to_model=args.path_talos_data.expanduser(), reduced=False) @@ -71,13 +101,14 @@ def main(): # Initialize simulator simulator = Simulator( - dt=dt, + dt=scen_params.dt, path_to_robot_urdf=args.path_talos_data.expanduser() / "talos_data" / "urdf" / "talos_full.urdf", model=talos, launch_gui=args.launch_gui, + n_solver_iter=scen_params.n_solver_iter, ) # Compute the right and left foot position as well as the base initial position @@ -110,7 +141,7 @@ def main(): w_com=10.0, w_mf=100.0, mu=1e-4, - dt=dt, + dt=scen_params.dt, locked_joints=list(talos.get_locked_joints_idx()), ) q_des, dq = solve_inverse_kinematics( @@ -132,10 +163,13 @@ def main(): # Then we start a stabilization phase where the robot has to maintain its CoM at a fixed position pb.setGravity(0, 0, -9.81) + # Setup the controller matrices + ctrler_mat = compute_preview_control_matrices(ctrler_params, scen_params.dt) + x_k = np.array([0.0, com_initial_target[0], 0.0, 0.0], dtype=float) y_k = np.array([0.0, com_initial_target[1], 0.0, 0.0], dtype=float) - for _ in range(math.ceil(1.0 / dt)): + for _ in range(math.ceil(1.0 / scen_params.dt)): # Get the current configuration of the robot from the simulator q_init = simulator.get_q(talos.model.nq) @@ -172,31 +206,31 @@ def main(): steps_pose, steps_ids = compute_steps_sequence( rf_initial_pose=rf_initial_pose, lf_initial_pose=lf_initial_pose, - n_steps=n_steps, - l_stride=l_stride, + n_steps=scen_params.n_steps, + l_stride=scen_params.l_stride, ) t, lf_path, rf_path, phases = compute_feet_trajectories( rf_initial_pose=rf_initial_pose, lf_initial_pose=lf_initial_pose, - n_steps=n_steps, + n_steps=scen_params.n_steps, steps_pose=steps_pose, - t_ss=t_ss, - t_ds=t_ds, - t_init=t_init, - t_end=t_end, - dt=dt, - traj_generator=BezierCurveFootPathGenerator(max_height_foot), + t_ss=scen_params.t_ss, + t_ds=scen_params.t_ds, + t_init=scen_params.t_init, + t_final=scen_params.t_end, + dt=scen_params.dt, + traj_generator=BezierCurveFootPathGenerator(scen_params.max_height_foot), ) zmp_ref = compute_zmp_ref( t=t, com_initial_pose=com_initial_target[0:2], steps=steps_pose[:, 0:2], - ss_t=t_ss, - ds_t=t_ds, - t_init=t_init, - t_final=t_end, + ss_t=scen_params.t_ss, + ds_t=scen_params.t_ds, + t_init=scen_params.t_init, + t_final=scen_params.t_end, interp_fn=cubic_spline_interpolation, ) @@ -221,6 +255,9 @@ def main(): rf_forces = np.zeros((len(phases), 1)) lf_forces = np.zeros((len(phases), 1)) + if args.record_video: + simulator.start_video_record() + # We start the walking phase for k, _ in enumerate(phases[:-2]): # Get the current configuration of the robot from the simulator @@ -256,8 +293,8 @@ def main(): params=ik_sol_params, ) - oMf_lf_tgt = pin.SE3(oMf_lf_tgt.rotation, lf_path[k + 1]) - + if phases[k + 1] > 0.0: + oMf_lf_tgt = pin.SE3(oMf_lf_tgt.rotation, lf_path[k + 1]) else: ik_sol_params.fixed_foot_frame = talos.left_foot_id ik_sol_params.moving_foot_frame = talos.right_foot_id @@ -272,7 +309,8 @@ def main(): params=ik_sol_params, ) - oMf_rf_tgt = pin.SE3(oMf_rf_tgt.rotation, rf_path[k + 1]) + if phases[k + 1] < 0.0: + oMf_rf_tgt = pin.SE3(oMf_rf_tgt.rotation, rf_path[k + 1]) simulator.apply_joints_pos_to_robot(q_des) @@ -300,6 +338,9 @@ def main(): rf_forces[k], lf_forces[k] = simulator.get_contact_forces() + if args.record_video: + simulator.stop_video_record(duration=t[-1]) + if args.plot_results: zmp_ref_plot = np.zeros((zmp_ref.shape[0], 3)) @@ -321,7 +362,7 @@ def main(): zmp_ref=zmp_ref_plot, ) - plot_contact_forces_and_state(t=t, force_rf=rf_forces, force_lf=lf_forces) + plot_contact_forces(t=t, force_rf=rf_forces, force_lf=lf_forces) plt.show() diff --git a/examples/example_5_walking_controller.py b/examples/example_5_walking_controller.py deleted file mode 100644 index 10e34c5..0000000 --- a/examples/example_5_walking_controller.py +++ /dev/null @@ -1,423 +0,0 @@ -import math -import argparse -from pathlib import Path - -import numpy as np -import pybullet as pb -import pinocchio as pin -from matplotlib import pyplot as plt - -from biped_walking_controller.foot import ( - compute_steps_sequence, - compute_swing_foot_pose, - compute_time_vector, - BezierCurveFootPathGenerator, -) - -from biped_walking_controller.inverse_kinematic import InvKinSolverParams, solve_inverse_kinematics -from biped_walking_controller.plot import plot_feet_and_com, plot_contact_forces_and_state - -from biped_walking_controller.preview_control import ( - PreviewControllerParams, - compute_preview_control_matrices, - update_control, - compute_zmp_ref, - cubic_spline_interpolation, -) - -from biped_walking_controller.model import Talos, q_from_base_and_joints - -from biped_walking_controller.simulation import ( - _snap_feet_to_plane, - _compute_base_from_foot_target, - Simulator, -) -from biped_walking_controller.state_machine import ( - WalkingStateMachineParams, - WalkingStateMachine, - WalkingState, - Foot, -) - - -def main(): - p = argparse.ArgumentParser() - p.add_argument("--path-talos-data", type=Path, help="Path to talos_data root") - p.add_argument("--plot-results", action="store_true") - p.add_argument("--launch-gui", action="store_true") - args = p.parse_args() - - np.set_printoptions(suppress=True, precision=3) - - dt = 1.0 / 240.0 - - # ZMP reference parameters - t_ss = 0.8 # Single support phase time window - t_ds = 0.3 # Double support phase time window - t_init = 2.0 # Initialization phase (transition from still position to first step) - t_end = 0.4 - n_steps = 5 # Number of steps executed by the robot - l_stride = 0.1 # Length of the stride - max_height_foot = 0.01 # Maximal height of the swing foot - - # Preview controller parameters - t_preview = 1.6 # Time horizon used for the preview controller - ctrler_params = PreviewControllerParams( - zc=0.89, - g=9.81, - Qe=1.0, - Qx=np.zeros((3, 3)), - R=1e-6, - n_preview_steps=int(round(t_preview / dt)), - ) - ctrler_mat = compute_preview_control_matrices(ctrler_params, dt) - - # Initialize Talos pinocchio model - talos = Talos(path_to_model=args.path_talos_data.expanduser(), reduced=False) - q_init = talos.set_and_get_default_pose() - - # Initialize simulator - simulator = Simulator( - dt=dt, - path_to_robot_urdf=args.path_talos_data.expanduser() - / "talos_data" - / "urdf" - / "talos_full.urdf", - model=talos, - launch_gui=args.launch_gui, - ) - - # Compute the right and left foot position as well as the base initial position - oMf_rf0 = talos.data.oMf[talos.right_foot_id].copy() - oMf_lf0 = talos.data.oMf[talos.left_foot_id].copy() - oMf_lf_tgt, oMf_rf_tgt = _snap_feet_to_plane(oMf_lf0, oMf_rf0) - - oMf_torso = talos.data.oMf[talos.torso_id].copy() - oMb_init = _compute_base_from_foot_target( - talos.model, talos.data, q_init, talos.left_foot_id, oMf_lf_tgt - ) - - # Compute the required initial configuration and apply them to the kinematic model - q_init = q_from_base_and_joints(q_init, oMb_init) - pin.forwardKinematics(talos.model, talos.data, q_init) - pin.updateFramePlacements(talos.model, talos.data) - - # set initial CoM target centered between feet at height zc - feet_mid = 0.5 * (oMf_lf_tgt.translation + oMf_rf_tgt.translation) - com_initial_target = np.array([feet_mid[0], feet_mid[1], ctrler_params.zc]) - - # We run a single inverse kinematic iteration to get the desired initial position of the robot - ik_sol_params = InvKinSolverParams( - fixed_foot_frame=talos.left_foot_id, - moving_foot_frame=talos.right_foot_id, - torso_frame=talos.torso_id, - model=talos.model, - data=talos.data, - w_torso=10.0, - w_com=10.0, - w_mf=100.0, - mu=1e-4, - dt=dt, - locked_joints=list(talos.get_locked_joints_idx()), - ) - q_des, dq = solve_inverse_kinematics( - q_init, com_initial_target, oMf_lf_tgt, oMf_rf_tgt, oMf_torso, ik_sol_params - ) - q_init = q_des - - # Update kinematic model and simulator - pin.forwardKinematics(talos.model, talos.data, q_init) - pin.updateFramePlacements(talos.model, talos.data) - simulator.reset_robot_configuration(q_init) - - # First we hard reset the position of the robot and let the simulation run for a few iterations with 0 gravity to - # stabilize the robot - pb.setGravity(0, 0, 0) - for _ in range(5): - simulator.step() # settle contacts before enabling motors - - # Then we start a stabilization phase where the robot has to maintain its CoM at a fixed position - pb.setGravity(0, 0, -9.81) - - x_k = np.array([0.0, com_initial_target[0], 0.0, 0.0], dtype=float) - y_k = np.array([0.0, com_initial_target[1], 0.0, 0.0], dtype=float) - - for _ in range(math.ceil(1.0 / dt)): - # Get the current configuration of the robot from the simulator - q_init = simulator.get_q(talos.model.nq) - - # Apply the configuration to the kinematic model - pin.forwardKinematics(talos.model, talos.data, q_init) - pin.updateFramePlacements(talos.model, talos.data) - - # # Get zmp ref horizon - zmp_ref_horizon = np.ones((ctrler_params.n_preview_steps - 1, 2)) * feet_mid[0:2] - - _, x_k, y_k = update_control( - ctrler_mat, feet_mid[0:2], zmp_ref_horizon, x_k.copy(), y_k.copy() - ) - - # The CoM target is meant to follow the computed x and y and stay at constant height zc from the feet - com_target = np.array([x_k[1], y_k[1], ctrler_params.zc]) - - # Stabilize at the position - q_des, dq = solve_inverse_kinematics( - q_init, com_target, oMf_lf_tgt, oMf_rf_tgt, oMf_torso, ik_sol_params - ) - - # Uncomment to follow the center of mass of the robot - simulator.update_camera_to_follow_pos(x_k[1], 0.0, 0.0) - - # This step is only here to start with the right initial position. Therefore we perform a reset - simulator.reset_robot_configuration(q_des) - simulator.step() - - lf_initial_pose = oMf_lf_tgt.translation - rf_initial_pose = oMf_rf_tgt.translation - - # Build ZMP reference to track - steps_pose, steps_ids = compute_steps_sequence( - rf_initial_pose=rf_initial_pose, - lf_initial_pose=lf_initial_pose, - n_steps=n_steps, - l_stride=l_stride, - ) - - t = compute_time_vector(t_ss, t_ds, t_init, t_end, n_steps, dt) - - zmp_ref = compute_zmp_ref( - t=t, - com_initial_pose=com_initial_target[0:2], - steps=steps_pose[:, 0:2], - ss_t=t_ss, - ds_t=t_ds, - t_init=t_init, - t_final=t_end, - interp_fn=cubic_spline_interpolation, - ) - - params = WalkingStateMachineParams( - t_init=t_init, t_end=t_end, t_ss=t_ss, t_ds=t_ds, force_threshold=50 - ) - state_machine = WalkingStateMachine(params=params, initial_state=WalkingState.INIT) - state_machine.update_steps(steps_pose, steps_ids) - - zmp_padded = np.vstack( - [zmp_ref, np.repeat(zmp_ref[-1][None, :], ctrler_params.n_preview_steps, axis=0)] - ) - - com_pin_pos = np.zeros((len(t), 3)) - com_ref_pos = np.zeros((len(t), 3)) - com_pb_pos = np.zeros((len(t), 3)) - - lf_ref_pos = np.zeros((len(t), 3)) - lf_pin_pos = np.zeros((len(t), 3)) - lf_pb_pos = np.zeros((len(t), 3)) - - rf_ref_pos = np.zeros((len(t), 3)) - rf_pin_pos = np.zeros((len(t), 3)) - rf_pb_pos = np.zeros((len(t), 3)) - - zmp_pos = np.zeros((len(t), 3)) - - rf_forces = np.zeros((len(t), 1)) - lf_forces = np.zeros((len(t), 1)) - states = np.zeros((len(t), 1)) - - foot_path_generator = BezierCurveFootPathGenerator(max_height_foot) - touchdown_vel = 0.0 - - # We start the walking phase - for k, _ in enumerate(t[:-2]): - # Get the current configuration of the robot from the simulator - q_init = simulator.get_q(talos.model.nq) - - # Apply the configuration to the kinematic model - pin.forwardKinematics(talos.model, talos.data, q_init) - pin.updateFramePlacements(talos.model, talos.data) - - # Update the control - zmp_ref_horizon = zmp_padded[k + 1 : k + ctrler_params.n_preview_steps] - - _, x_k, y_k = update_control( - ctrler_mat, zmp_padded[k], zmp_ref_horizon, x_k.copy(), y_k.copy() - ) - - # The CoM target is meant to follow the computed x and y and stay at constant height zc from the feet - com_target = np.array([x_k[1], y_k[1], ctrler_params.zc]) - - # Get contact forces and update the walking state machine - rf_forces[k], lf_forces[k] = simulator.get_contact_forces() - state_machine.update(t=k * dt, rf_contact_force=rf_forces[k], lf_contact_force=lf_forces[k]) - states[k] = state_machine.get_current_state().value - - zmp_pos[k] = simulator.get_zmp_pose() - - step_idx = state_machine.get_step_idx() - - # Alternate between feet - if state_machine.get_current_state() == WalkingState.SS_RIGHT: - ik_sol_params.fixed_foot_frame = talos.right_foot_id - ik_sol_params.moving_foot_frame = talos.left_foot_id - - if step_idx < 2: - lf_start = lf_initial_pose - lf_target = steps_pose[step_idx + 1] - else: - lf_start = steps_pose[step_idx - 1] - lf_target = steps_pose[step_idx + 1] - - rf_pose = steps_pose[step_idx] - lf_pose = compute_swing_foot_pose( - t_state=state_machine.get_elapsed_time_in_state(k * dt), - params=params, - step_start=lf_start, - step_target=lf_target, - touchdown_extension_vel=touchdown_vel, - path_generator=foot_path_generator, - ) - - oMf_lf = pin.SE3(oMf_lf_tgt.rotation, lf_pose) - q_des, dq = solve_inverse_kinematics( - q_init, - com_target, - oMf_fixed_foot=pin.SE3(oMf_rf_tgt.rotation, rf_pose), - oMf_moving_foot=oMf_lf, - oMf_torso=oMf_torso, - params=ik_sol_params, - ) - elif state_machine.get_current_state() == WalkingState.SS_LEFT: - if step_idx < 2: - rf_start = rf_initial_pose - rf_target = steps_pose[step_idx + 1] - else: - rf_start = steps_pose[step_idx - 1] - rf_target = steps_pose[step_idx + 1] - - ik_sol_params.fixed_foot_frame = talos.left_foot_id - ik_sol_params.moving_foot_frame = talos.right_foot_id - - lf_pose = steps_pose[step_idx] - rf_pose = compute_swing_foot_pose( - t_state=state_machine.get_elapsed_time_in_state(k * dt), - params=params, - step_start=rf_start, - step_target=rf_target, - touchdown_extension_vel=touchdown_vel, - path_generator=foot_path_generator, - ) - - oMf_rf = pin.SE3(oMf_rf_tgt.rotation, rf_pose) - q_des, dq = solve_inverse_kinematics( - q_init, - com_target, - oMf_fixed_foot=pin.SE3(oMf_lf_tgt.rotation, lf_pose), - oMf_moving_foot=oMf_rf, - oMf_torso=oMf_torso, - params=ik_sol_params, - ) - elif state_machine.get_current_state() == WalkingState.END: - if steps_ids[-1] == Foot.RIGHT: - rf_pose = steps_pose[-2] - lf_pose = steps_pose[-1] - else: - rf_pose = steps_pose[-1] - lf_pose = steps_pose[-2] - - ik_sol_params.fixed_foot_frame = talos.left_foot_id - ik_sol_params.moving_foot_frame = talos.right_foot_id - - q_des, dq = solve_inverse_kinematics( - q_init, - com_target, - oMf_fixed_foot=pin.SE3(oMf_lf_tgt.rotation, lf_pose), - oMf_moving_foot=pin.SE3(oMf_rf_tgt.rotation, rf_pose), - oMf_torso=oMf_torso, - params=ik_sol_params, - ) - else: - if steps_ids[step_idx] == Foot.RIGHT: - if step_idx == 0: - rf_pose = steps_pose[step_idx] - lf_pose = lf_initial_pose - else: - rf_pose = steps_pose[step_idx] - lf_pose = steps_pose[step_idx - 1] - else: - if step_idx == 0: - rf_pose = rf_initial_pose - lf_pose = steps_pose[step_idx] - else: - rf_pose = steps_pose[step_idx - 1] - lf_pose = steps_pose[step_idx] - - ik_sol_params.fixed_foot_frame = talos.left_foot_id - ik_sol_params.moving_foot_frame = talos.right_foot_id - - q_des, dq = solve_inverse_kinematics( - q_init, - com_target, - oMf_fixed_foot=pin.SE3(oMf_lf_tgt.rotation, lf_pose), - oMf_moving_foot=pin.SE3(oMf_rf_tgt.rotation, rf_pose), - oMf_torso=oMf_torso, - params=ik_sol_params, - ) - - simulator.apply_joints_pos_to_robot(q_des) - - # Uncomment to follow the center of mass of the robot - simulator.update_camera_to_follow_pos(x_k[1], 0.0, 0.0) - - simulator.step() - - if args.plot_results: - pin.computeCentroidalMap(talos.model, talos.data, q_init) - com_pin = pin.centerOfMass(talos.model, talos.data, q_init) - - # Store position of CoM, left and right feet - com_ref_pos[k] = com_target - com_pin_pos[k] = com_pin - com_pb_pos[k] = simulator.get_robot_com_position() - - lf_ref_pos[k] = lf_pose - lf_pin_pos[k] = talos.data.oMf[talos.left_foot_id].translation - lf_pb_pos[k], _ = simulator.get_robot_frame_pos("leg_left_6_link") - - rf_ref_pos[k] = rf_pose - rf_pin_pos[k] = talos.data.oMf[talos.right_foot_id].translation - rf_pb_pos[k], _ = simulator.get_robot_frame_pos("leg_right_6_link") - - rf_forces[k], lf_forces[k] = simulator.get_contact_forces() - - if args.plot_results: - - zmp_ref_plot = np.zeros((zmp_ref.shape[0], 3)) - zmp_ref_plot[:, :2] = zmp_ref - - plot_feet_and_com( - title_prefix="Walking controller", - t=t, - lf_pin_pos=lf_pin_pos, - rf_pin_pos=rf_pin_pos, - lf_ref_pos=lf_ref_pos, - rf_ref_pos=rf_ref_pos, - lf_pb_pos=lf_pb_pos, - rf_pb_pos=rf_pb_pos, - com_ref_pos=com_ref_pos, - com_pb_pos=com_pb_pos, - com_pin_pos=com_pin_pos, - zmp_pb=zmp_pos, - zmp_ref=zmp_ref_plot, - ) - - plot_contact_forces_and_state(t=t, force_rf=rf_forces, force_lf=lf_forces, states=states) - - plt.show() - - # Infinite loop to display the ending position - while True: - simulator.step() - - -if __name__ == "__main__": - main()