From 2034f139dff1243430abe2d9b6a3932f736662ec Mon Sep 17 00:00:00 2001 From: Yifei Zhou <75226267+yf-zhou@users.noreply.github.com> Date: Sat, 16 Nov 2024 17:46:50 -0500 Subject: [PATCH] 891 vis walking metrics (#902) * visualize imu and position * small changes * small changes * reshape footprint rectangles * clean up pt 3??? * remove if statements - graph everything --- .gitignore | 3 + .../src/soccer_pycontrol/model/sensors.py | 4 +- .../soccer_pycontrol/walk_engine/navigator.py | 193 +++++++++++++++++- .../soccer_pycontrol/test/test_placo.py | 14 +- 4 files changed, 205 insertions(+), 9 deletions(-) diff --git a/.gitignore b/.gitignore index 2e9384efd..01242bf08 100644 --- a/.gitignore +++ b/.gitignore @@ -34,3 +34,6 @@ venv/ .mypy_cache/ cuda-ubuntu2004.pin.1 *.webm + +# vscode +.vscode/* diff --git a/soccer_control/soccer_pycontrol/src/soccer_pycontrol/model/sensors.py b/soccer_control/soccer_pycontrol/src/soccer_pycontrol/model/sensors.py index fe996b9ed..d9ec04eb6 100644 --- a/soccer_control/soccer_pycontrol/src/soccer_pycontrol/model/sensors.py +++ b/soccer_control/soccer_pycontrol/src/soccer_pycontrol/model/sensors.py @@ -19,13 +19,13 @@ def __init__(self, body: pb.loadURDF): self.imu_ready = False self.get_imu() # to init - def get_pose(self) -> Transformation: + def get_pose(self, link=None) -> Transformation: """ Get the 3D pose of the robot :return: The 3D pose of the robot """ - [position, quaternion] = pb.getLinkState(self.body, linkIndex=self.imu_link)[4:6] # TODO double check + [position, quaternion] = pb.getLinkState(self.body, linkIndex=self.imu_link if not link else link)[4:6] # TODO double check return Transformation(position=position, quaternion=quaternion) def get_imu(self) -> Transformation: diff --git a/soccer_control/soccer_pycontrol/src/soccer_pycontrol/walk_engine/navigator.py b/soccer_control/soccer_pycontrol/src/soccer_pycontrol/walk_engine/navigator.py index 63b4331bb..8f4cc8fe2 100644 --- a/soccer_control/soccer_pycontrol/src/soccer_pycontrol/walk_engine/navigator.py +++ b/soccer_control/soccer_pycontrol/src/soccer_pycontrol/walk_engine/navigator.py @@ -1,7 +1,9 @@ import math import time +from collections import defaultdict from typing import List, Union +import matplotlib.pyplot as plt import numpy as np import scipy from soccer_pycontrol.model.bez import Bez @@ -14,7 +16,7 @@ # TODO change to trajectory controller class Navigator: - def __init__(self, world: PybulletWorld, bez: Bez, imu_feedback_enabled: bool = False): + def __init__(self, world: PybulletWorld, bez: Bez, imu_feedback_enabled: bool = False, record_walking_metrics: bool = True): self.world = world self.bez = bez self.imu_feedback_enabled = imu_feedback_enabled @@ -48,8 +50,16 @@ def __init__(self, world: PybulletWorld, bez: Bez, imu_feedback_enabled: bool = self.error_tol = 0.01 # in m TODO add as a param and in the ros version + # joints + self.left_ankle_index = self.bez.motor_control.motor_names.index("left_ankle_roll") + self.right_ankle_index = self.bez.motor_control.motor_names.index("right_ankle_roll") + # self.torso_index = self.bez.motor_control.body. + + self.record_walking_metrics = record_walking_metrics + self.walking_data = defaultdict(list) + # TODO could make input a vector - def walk(self, target_goal: Union[Transformation, List], ball_mode: bool = False): + def walk(self, target_goal: Union[Transformation, List], ball_mode: bool = False, display_metrics: bool = False): if isinstance(target_goal, Transformation): if ball_mode: self.walk_ball(target_goal) @@ -58,6 +68,8 @@ def walk(self, target_goal: Union[Transformation, List], ball_mode: bool = False elif isinstance(target_goal, list): # [d_x: float = 0.0, d_y: float = 0.0, d_theta: float = 0.0, nb_steps: int = 10, t_goal: float = 10] self.walk_time(target_goal) + if self.record_walking_metrics and display_metrics: + self.display_walking_metrics(show_targets=isinstance(target_goal, Transformation)) self.ready() def find_new_vel(self, goal_loc: list, curr_loc: list = (0, 0)): @@ -203,6 +215,9 @@ def walk_loop(self, t: float): # self.foot_step_planner.robot.set_T_world_fbase(T) # self.foot_step_planner.robot.update_kinematics() + if self.record_walking_metrics: + self.update_walking_metrics(t) + return t @staticmethod @@ -253,6 +268,180 @@ def filter_joints(self) -> List[int]: joints[self.bez.motor_control.motor_names.index(joint)] = self.foot_step_planner.robot.get_joint(joint) return joints + def display_walking_metrics(self, show_targets: bool = False) -> None: + fig, (ax_imu0, ax_imu1, ax_imu2) = plt.subplots(3, 1, sharex=True) + fig.canvas.set_window_title("imu") + + imu_0 = np.array(np.array(self.walking_data["IMU_0"]).transpose()) + ax_imu0.plot(imu_0[0, :], imu_0[1, :]) + if show_targets: + ax_imu0.plot( + imu_0[0, :], + np.ones(imu_0[0, :].shape) * self.nav_yaw_pid.setpoint, + linewidth=0.5, + color="r", + label=f"target yaw ({self.nav_yaw_pid.setpoint})", + ) + ax_imu0.set_title("yaw") + ax_imu0.grid() + + imu_1 = np.array(np.array(self.walking_data["IMU_1"]).transpose()) + ax_imu1.plot(imu_1[0, :], imu_1[1, :]) + ax_imu1.plot(imu_1[0, :], np.zeros(imu_1[0, :].shape), linewidth=0.5, color="r", label=f"target pitch ({0.0})") + ax_imu1.set_title("pitch") + ax_imu1.grid() + + imu_2 = np.array(np.array(self.walking_data["IMU_2"]).transpose()) + ax_imu2.plot(imu_2[0, :], imu_2[1, :]) + ax_imu2.plot(imu_2[0, :], np.zeros(imu_2[0, :].shape), linewidth=0.5, color="r", label=f"target roll ({0.0})") + ax_imu2.set_title("roll") + ax_imu2.grid() + + plt.subplots_adjust(wspace=0.3, hspace=0.5) + + plt.show() + + # fig, ax = plt.subplots(3, 1) + fig = plt.figure(figsize=(5, 7)) + fig.canvas.set_window_title("position") + gs = fig.add_gridspec(10, 1) + + ax_position = fig.add_subplot(gs[:6]) + ax_pos_err = fig.add_subplot(gs[7]) + ax_yaw_err = fig.add_subplot(gs[9]) + + target_x = 0 if not show_targets else self.nav_x_pid.setpoint + target_y = 0 if not show_targets else self.nav_y_pid.setpoint + target_yaw = 0 if not show_targets else self.nav_yaw_pid.setpoint + + position = np.array(self.walking_data["POSITION"]).transpose() + ax_position.plot(position[1, :], position[2, :]) + ax_position.plot(position[1, 0], position[2, 0], "yo", label="start point") + ax_position.plot(position[1, -1], position[2, -1], "go", label="end point") + if show_targets: + ax_position.plot(self.nav_x_pid.setpoint, self.nav_y_pid.setpoint, "ro", label="target point") + ax_position.set_title("position") + ax_position.set_xlabel("x") + ax_position.set_ylabel("y") + ax_position.grid() + ax_position.legend() + + if show_targets: + ax_pos_err.plot(position[0, :], np.linalg.norm(position[1:3, :].transpose() - np.array([target_x, target_y]), axis=1)) + ax_pos_err.plot(position[0, :], np.zeros(position[0, :].shape), linewidth=0.5, color="r") + ax_pos_err.set_title("position error") + ax_pos_err.set_ylabel("euclidean distance") + ax_pos_err.grid() + + ax_yaw_err.plot(position[0, :], position[3, :] - target_yaw) + ax_yaw_err.plot(position[0, :], np.zeros(position[0, :].shape), linewidth=0.5, color="r") + ax_yaw_err.set_title("orientation error") + ax_yaw_err.grid() + + plt.show() + + fig = plt.figure() + fig.canvas.set_window_title("3d position") + gs = fig.add_gridspec(1, 1) + + ax_path = fig.add_subplot(gs[0], projection="3d") + + left = np.array(self.walking_data["LEFT_FOOT"]).T + right = np.array(self.walking_data["RIGHT_FOOT"]).T + com = np.array(self.walking_data["COM"]).T + + ax_path.plot(left[1, :], left[2, :], left[3, :], label="left foot") + ax_path.plot(right[1, :], right[2, :], right[3, :], label="right foot") + ax_path.plot(com[1, :], com[2, :], com[3, :], label="centre of mass") + # ax_path.plot(com[1, :], com[2, :], np.zeros(com[3, :].shape), label="centre of mass projected") + ax_path.set_title("robot stance") + ax_path.grid() + ax_path.legend() + + plt.show() + + fig = plt.figure() + fig.canvas.set_window_title("footprints") + gs = fig.add_gridspec(1, 1) + + ax_footprints = fig.add_subplot(gs[0]) + + # both feet on the ground + ground = 0.025 + left_grounded_i = np.where(left[3] < ground)[0] + right_grounded_i = np.where(right[3] < ground)[0] + + rectangle_width = 0.1 + rectangle_length = 0.05 + rectangle_x_offset = 0 + rectangle_y_offset = 0 + + make_rotation = lambda theta, x, y: np.array([[np.cos(theta), -np.sin(theta), x], [np.sin(theta), np.cos(theta), y]]) + rectangle = np.array( + [ + [rectangle_x_offset - rectangle_width / 2, rectangle_y_offset + rectangle_length / 2, 1], + [rectangle_x_offset + rectangle_width / 2, rectangle_y_offset + rectangle_length / 2, 1], + [rectangle_x_offset + rectangle_width / 2, rectangle_y_offset - rectangle_length / 2, 1], + [rectangle_x_offset - rectangle_width / 2, rectangle_y_offset - rectangle_length / 2, 1], + [rectangle_x_offset - rectangle_width / 2, rectangle_y_offset + rectangle_length / 2, 1], + ] + ) + make_rectangle = lambda theta, x, y: np.array([make_rotation(theta, x, y) @ rectangle[i].T for i in range(5)]) + + yaw_data = np.array(self.walking_data["IMU_0"]).T + + left_rectangles = [] + for i in left_grounded_i: + if np.any(np.where(yaw_data[0] == left[0, i])[0]): + left_rectangles.append(make_rectangle(yaw_data[1][np.where(yaw_data[0] == left[0, i])[0][0]], left[1, i], left[2, i]).T) + + right_rectangles = [] + for i in right_grounded_i: + if np.any(np.where(yaw_data[0] == right[0, i])[0]): + right_rectangles.append(make_rectangle(yaw_data[1][np.where(yaw_data[0] == right[0, i])[0][0]], right[1, i], right[2, i]).T) + + for r in left_rectangles: + ax_footprints.plot(r[0], r[1]) + for r in right_rectangles: + ax_footprints.plot(r[0], r[1]) + + ax_footprints.plot(com[1, :], com[2, :], label="centre of mass") + + ax_footprints.set_aspect("equal") + + plt.show() + + def clear_walking_metrics(self, target_data: list = None) -> None: + """reinitialize walking data""" + if not target_data: + # clear all + self.walking_data = defaultdict(list) + return + for name in self.walking_data.keys(): + if name in target_data: + self.walking_data[name] = [] + + def update_walking_metrics(self, t: float) -> None: + """update stored data for time t""" + + # IMU data + imu_data = self.bez.sensors.get_imu() + self.walking_data["IMU_0"].append((t, imu_data[0])) + self.walking_data["IMU_1"].append((t, imu_data[1])) + self.walking_data["IMU_2"].append((t, imu_data[2])) + + # position data + pose = self.bez.sensors.get_pose() + self.walking_data["POSITION"].append((t, pose.position[0], pose.position[1], pose.orientation_euler[0])) + + # feet and COM data + left_foot = self.bez.sensors.get_pose(self.left_ankle_index) # 13 for left foot joint + right_foot = self.bez.sensors.get_pose(self.right_ankle_index) # 7 for right foot joint + com = self.bez.sensors.get_pose(1) # 1 for torso joint + self.walking_data["LEFT_FOOT"].append((t, left_foot.position[0], left_foot.position[1], left_foot.position[2])) + self.walking_data["RIGHT_FOOT"].append((t, right_foot.position[0], right_foot.position[1], right_foot.position[2])) + self.walking_data["COM"].append((t, com.position[0], com.position[1], com.position[2])) + if __name__ == "__main__": world = PybulletWorld( diff --git a/soccer_control/soccer_pycontrol/test/test_placo.py b/soccer_control/soccer_pycontrol/test/test_placo.py index 592b3188d..10e9fb46f 100644 --- a/soccer_control/soccer_pycontrol/test/test_placo.py +++ b/soccer_control/soccer_pycontrol/test/test_placo.py @@ -6,7 +6,7 @@ from soccer_common import Transformation -REAL_TIME = True +REAL_TIME = False class TestPlaco(unittest.TestCase): @@ -21,11 +21,15 @@ def test_bez1(self): real_time=REAL_TIME, rate=200, ) - self.bez = Bez(robot_model="bez2", pose=Transformation()) - walk = Navigator(self.world, self.bez) - target_goal = [0.08, 0, 0, 3, 500] + self.bez = Bez(robot_model="bez1", pose=Transformation()) + walk = Navigator(self.world, self.bez, imu_feedback_enabled=True) + # walk.ready() + # self.bez.motor_control.set_motor() + # walk.wait(50) + target_goal = [0.08, 0, 0, 3, 10] # target_goal = Transformation(position=[1, 0, 0], euler=[0, 0, 0]) - walk.walk(target_goal) + print("STARTING WALK") + walk.walk(target_goal, display_metrics=True) # walk.wait(1000) # TODO fix