diff --git a/README.md b/README.md index 169fd56..8df8e30 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,17 @@ pip install --no-deps -e . ## Usage Take a look at the [examples](./examples) folder! +To run the examples you will need to install notebook +``` +conda install notebook +``` + +> [!NOTE] +> Running the example using the [drake simulator](https://drake.mit.edu/) needs installing the following additional dependencies, +``` +conda install meshio tqdm +pip install drake git+https://github.com/akhilsathuluri/odio_urdf.git@comodo +``` ### Maintainer diff --git a/examples/drake_walking.ipynb b/examples/drake_walking.ipynb new file mode 100644 index 0000000..7fa1a2b --- /dev/null +++ b/examples/drake_walking.ipynb @@ -0,0 +1,334 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Comodo Exmaple \n", + "This examples, load a basic robot model (i.e. composed only of basic shapes), modifies the links of such a robot model by elongating the legs, define instances of the TSID (Task Based Inverse Dynamics) and Centroidal MPC controller and simulate the behavior of the robot using drake. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Comodo import\n", + "from comodo.drakeSimulator.drakeSimulator import DrakeSimulator\n", + "from comodo.robotModel.robotModel import RobotModel\n", + "from comodo.robotModel.createUrdf import createUrdf\n", + "from comodo.centroidalMPC.centroidalMPC import CentroidalMPC\n", + "from comodo.centroidalMPC.mpcParameterTuning import MPCParameterTuning\n", + "from comodo.TSIDController.TSIDParameterTuning import TSIDParameterTuning\n", + "from comodo.TSIDController.TSIDController import TSIDController" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# General import \n", + "import xml.etree.ElementTree as ET\n", + "import numpy as np\n", + "import tempfile\n", + "import urllib.request\n", + "import logging" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# set the logging level\n", + "logging.getLogger().setLevel(logging.INFO)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Getting stickbot urdf file and convert it to string \n", + "urdf_robot_file = tempfile.NamedTemporaryFile(mode=\"w+\")\n", + "url = 'https://raw.githubusercontent.com/icub-tech-iit/ergocub-gazebo-simulations/master/models/stickBot/model.urdf'\n", + "urllib.request.urlretrieve(url, urdf_robot_file.name)\n", + "# Load the URDF file\n", + "tree = ET.parse(urdf_robot_file.name)\n", + "root = tree.getroot()\n", + "\n", + "# Convert the XML tree to a string\n", + "robot_urdf_string_original = ET.tostring(root)\n", + "\n", + "create_urdf_instance = createUrdf(\n", + " original_urdf_path=urdf_robot_file.name, save_gazebo_plugin=False\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Define parametric links and controlled joints \n", + "legs_link_names = [\"hip_3\", \"lower_leg\"]\n", + "joint_name_list = [\n", + " \"r_shoulder_pitch\",\n", + " \"r_shoulder_roll\",\n", + " \"r_shoulder_yaw\",\n", + " \"r_elbow\",\n", + " \"l_shoulder_pitch\",\n", + " \"l_shoulder_roll\",\n", + " \"l_shoulder_yaw\",\n", + " \"l_elbow\",\n", + " \"r_hip_pitch\",\n", + " \"r_hip_roll\",\n", + " \"r_hip_yaw\",\n", + " \"r_knee\",\n", + " \"r_ankle_pitch\",\n", + " \"r_ankle_roll\",\n", + " \"l_hip_pitch\",\n", + " \"l_hip_roll\",\n", + " \"l_hip_yaw\",\n", + " \"l_knee\",\n", + " \"l_ankle_pitch\",\n", + " \"l_ankle_roll\",\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Define the robot modifications\n", + "modifications = {}\n", + "for item in legs_link_names:\n", + " left_leg_item = \"l_\" + item\n", + " right_leg_item = \"r_\" + item\n", + " modifications.update({left_leg_item: 1.2})\n", + " modifications.update({right_leg_item: 1.2})\n", + "# Motors Parameters \n", + "Im_arms = 1e-3*np.ones(4) # from 0-4\n", + "Im_legs = 1e-3*np.ones(6) # from 5-10\n", + "kv_arms = 0.001*np.ones(4) # from 11-14\n", + "kv_legs = 0.001*np.ones(6) #from 20\n", + "\n", + "Im = np.concatenate((Im_arms, Im_arms, Im_legs, Im_legs))\n", + "kv = np.concatenate((kv_arms, kv_arms, kv_legs, kv_legs))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Modify the robot model and initialize\n", + "create_urdf_instance.modify_lengths(modifications)\n", + "urdf_robot_string = create_urdf_instance.write_urdf_to_file()\n", + "create_urdf_instance.reset_modifications()\n", + "robot_model_init = RobotModel(urdf_robot_string, \"stickBot\", joint_name_list)\n", + "s_des, xyz_rpy, H_b = robot_model_init.compute_desired_position_walking()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Define simulator and initial position\n", + "drake_instance = DrakeSimulator()\n", + "drake_instance.set_visualize_robot_flag(True)\n", + "drake_instance.load_model(robot_model_init, s=s_des, xyz_rpy=xyz_rpy, kv_motors=kv, Im=Im)\n", + "s, ds, tau = drake_instance.get_state()\n", + "t = drake_instance.get_simulation_time()\n", + "H_b = drake_instance.get_base()\n", + "# beware of the velocity ordering\n", + "w_b = drake_instance.get_base_velocity()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Define the controller parameters and instantiate the controller\n", + "# Controller Parameters\n", + "tsid_parameter = TSIDParameterTuning()\n", + "mpc_parameters = MPCParameterTuning()\n", + "\n", + "# TSID Instance\n", + "TSID_controller_instance = TSIDController(frequency=0.01, robot_model=robot_model_init)\n", + "TSID_controller_instance.define_tasks(tsid_parameter)\n", + "TSID_controller_instance.set_state_with_base(s, ds, H_b, w_b, t)\n", + "\n", + "# MPC Instance\n", + "step_lenght = 0.1\n", + "mpc = CentroidalMPC(robot_model=robot_model_init, step_length=step_lenght)\n", + "mpc.intialize_mpc(mpc_parameters=mpc_parameters)\n", + "\n", + "# Set desired quantities\n", + "mpc.configure(s_init=s_des, H_b_init=H_b)\n", + "TSID_controller_instance.compute_com_position()\n", + "mpc.define_test_com_traj(TSID_controller_instance.COM.toNumPy())\n", + "\n", + "# TODO: Set initial robot state and plan trajectories \n", + "drake_instance.step(1)\n", + "\n", + "# Reading the state \n", + "s, ds, tau = drake_instance.get_state()\n", + "H_b = drake_instance.get_base()\n", + "w_b = drake_instance.get_base_velocity()\n", + "t = drake_instance.get_simulation_time()\n", + "\n", + "# MPC\n", + "mpc.set_state_with_base(s=s, s_dot=ds, H_b=H_b, w_b=w_b, t=t)\n", + "mpc.initialize_centroidal_integrator(s=s, s_dot=ds, H_b=H_b, w_b=w_b, t=t)\n", + "mpc_output = mpc.plan_trajectory()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Set loop variables \n", + "TIME_TH = 20.0\n", + "\n", + "# Define number of steps\n", + "n_step = int(\n", + " TSID_controller_instance.frequency / drake_instance.get_simulation_frequency()\n", + ")\n", + "n_step_mpc_tsid = int(mpc.get_frequency_seconds() / TSID_controller_instance.frequency)\n", + "\n", + "counter = 0\n", + "mpc_success = True\n", + "energy_tot = 0.0\n", + "succeded_controller = True" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Simulation-control loop\n", + "if drake_instance.visualize_robot_flag:\n", + " drake_instance.meshcat.StartRecording()\n", + "while t < TIME_TH:\n", + " \n", + " # Reading robot state from simulator\n", + " s, ds, tau = drake_instance.get_state()\n", + " energy_i = np.linalg.norm(tau)\n", + " H_b = drake_instance.get_base()\n", + " w_b = drake_instance.get_base_velocity()\n", + " t = drake_instance.get_simulation_time()\n", + "\n", + " # Update TSID\n", + " TSID_controller_instance.set_state_with_base(s=s, s_dot=ds, H_b=H_b, w_b=w_b, t=t)\n", + "\n", + " # MPC plan \n", + " if counter == 0:\n", + " mpc.set_state_with_base(s=s, s_dot=ds, H_b=H_b, w_b=w_b, t=t)\n", + " mpc.update_references()\n", + " mpc_success = mpc.plan_trajectory()\n", + " mpc.contact_planner.advance_swing_foot_planner()\n", + " if not (mpc_success):\n", + " logging.error(\"MPC failed\")\n", + " break\n", + "\n", + " # Reading new references\n", + " com, dcom, forces_left, forces_right, ang_mom = mpc.get_references()\n", + " left_foot, right_foot = mpc.contact_planner.get_references_swing_foot_planner()\n", + "\n", + " # Update references TSID \n", + " TSID_controller_instance.update_task_references_mpc(\n", + " com=com,\n", + " dcom=dcom,\n", + " ddcom=np.zeros(3),\n", + " left_foot_desired=left_foot,\n", + " right_foot_desired=right_foot,\n", + " s_desired=np.array(s_des),\n", + " wrenches_left=forces_left,\n", + " wrenches_right=forces_right,\n", + " ang_mom_desired=ang_mom,\n", + " )\n", + "\n", + " # Run control \n", + " succeded_controller = TSID_controller_instance.run()\n", + "\n", + " if not (succeded_controller):\n", + " logging.error(\"Controller failed\")\n", + " break\n", + "\n", + " tau = TSID_controller_instance.get_torque()\n", + "\n", + " # Step the simulator\n", + " drake_instance.set_input(tau)\n", + " # drake_instance.step_with_motors(n_step=n_step, torque=tau)\n", + " try:\n", + " drake_instance.step(n_step=n_step)\n", + " except RuntimeError as e:\n", + " logging.error(e)\n", + " break\n", + " counter = counter + 1\n", + "\n", + " if counter == n_step_mpc_tsid:\n", + " counter = 0\n", + "if drake_instance.visualize_robot_flag:\n", + " drake_instance.meshcat.StopRecording()\n", + " drake_instance.meshcat.PublishRecording()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Closing visualization\n", + "drake_instance.close_visualization()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "walk", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.13" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/examples/mujoco_walking.ipynb b/examples/mujoco_walking.ipynb index cf3b743..a271288 100644 --- a/examples/mujoco_walking.ipynb +++ b/examples/mujoco_walking.ipynb @@ -141,7 +141,7 @@ "t = mujoco_instance.get_simulation_time()\n", "H_b = mujoco_instance.get_base()\n", "w_b = mujoco_instance.get_base_velocity()\n", - "mujoco_instance.set_visualize_robot_flag(False)" + "mujoco_instance.set_visualize_robot_flag(True)" ] }, { @@ -236,7 +236,7 @@ " break\n", "\n", " # Reading new references\n", - " com, dcom, forces_left, forces_right = mpc.get_references()\n", + " com, dcom, forces_left, forces_right, ang_mom = mpc.get_references()\n", " left_foot, right_foot = mpc.contact_planner.get_references_swing_foot_planner()\n", " \n", " # Update references TSID \n", @@ -249,6 +249,7 @@ " s_desired=np.array(s_des),\n", " wrenches_left=forces_left,\n", " wrenches_right=forces_right,\n", + " ang_mom_desired = ang_mom\n", " )\n", "\n", " # Run control \n", @@ -266,7 +267,9 @@ " counter = counter + 1\n", "\n", " if counter == n_step_mpc_tsid:\n", - " counter = 0" + " counter = 0\n", + "\n", + "mujoco_instance.close_visualization()" ] }, { @@ -278,6 +281,13 @@ "# Closing visualization\n", "mujoco_instance.close_visualization()" ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/src/comodo/TSIDController/TSIDController.py b/src/comodo/TSIDController/TSIDController.py index 72bca44..eafbd48 100644 --- a/src/comodo/TSIDController/TSIDController.py +++ b/src/comodo/TSIDController/TSIDController.py @@ -1,14 +1,14 @@ -from comodo.abstractClasses.controller import Controller -from comodo.TSIDController.TSIDParameterTuning import TSIDParameterTuning import bipedal_locomotion_framework as blf -import numpy as np import idyntree.bindings as iDynTree import manifpy +import numpy as np + +from comodo.abstractClasses.controller import Controller +from comodo.TSIDController.TSIDParameterTuning import TSIDParameterTuning class TSIDController(Controller): def __init__(self, frequency, robot_model) -> None: - blf.text_logging.set_verbosity(blf.text_logging.Verbosity.Debug) self.controller = blf.tsid.QPTSID() self.gravity = iDynTree.Vector3() @@ -44,7 +44,6 @@ def define_varible_handler(self): self.var_handler.add_variable(self.variable_name_right_contact, 6) def define_tasks(self, tsid_parameters: TSIDParameterTuning): - self.define_kyndyn() self.controller_variable_handler = blf.parameters_handler.StdParametersHandler() self.controller_variable_handler.set_parameter_string( @@ -303,6 +302,63 @@ def define_tasks(self, tsid_parameters: TSIDParameterTuning): self.root_link_task.set_kin_dyn(self.kindyn) self.root_link_task.initialize(param_handler=self.root_link_task_param_handler) + # feet regularization terms + self.left_foot_regularization_task = blf.tsid.VariableRegularizationTask() + self.left_foot_regularization_task_name = "left_foot_regularization_task" + self.left_foot_regularization_task_priority = 1 + self.left_foot_regularization_task_weight = 1e-1 * np.ones(6) + self.left_foot_regularization_task_param_handler = ( + blf.parameters_handler.StdParametersHandler() + ) + self.left_foot_regularization_task_param_handler.set_parameter_string( + name="variable_name", value=self.variable_name_left_contact + ) + self.left_foot_regularization_task_param_handler.set_parameter_int( + name="variable_size", value=6 + ) + self.left_foot_regularization_task.initialize( + param_handler=self.left_foot_regularization_task_param_handler + ) + + self.right_foot_regularization_task = blf.tsid.VariableRegularizationTask() + self.right_foot_regularization_task_name = "right_foot_regularization_task" + self.right_foot_regularization_task_priority = 1 + self.right_foot_regularization_task_weight = 1e-1 * np.ones(6) + self.right_foot_regularization_task_param_handler = ( + blf.parameters_handler.StdParametersHandler() + ) + self.right_foot_regularization_task_param_handler.set_parameter_string( + name="variable_name", value=self.variable_name_right_contact + ) + self.right_foot_regularization_task_param_handler.set_parameter_int( + name="variable_size", value=6 + ) + self.right_foot_regularization_task.initialize( + param_handler=self.right_foot_regularization_task_param_handler + ) + + ## Angular momentum task + self.angular_momentum_task = blf.tsid.AngularMomentumTask() + self.angular_momentum_task_name = "angular_momentum_task" + self.angular_momentum_task_priority = 1 + self.angular_momentum_task_weight = 1e0 * np.ones(3) + self.angular_momentum_task_parameter_handler = ( + blf.parameters_handler.StdParametersHandler() + ) + self.angular_momentum_task_parameter_handler.set_group( + "CONTACT_0", contact_group_left + ) + self.angular_momentum_task_parameter_handler.set_group( + "CONTACT_1", contact_group_right + ) + self.angular_momentum_task_parameter_handler.set_parameter_float( + name="kp", value=10.0 + ) + self.angular_momentum_task.set_kin_dyn(self.kindyn) + self.angular_momentum_task.initialize( + self.angular_momentum_task_parameter_handler + ) + ## Add tasks to the controller self.controller.add_task( self.CoM_Task, self.CoM_task_name, self.CoM_task_priority @@ -349,6 +405,28 @@ def define_tasks(self, tsid_parameters: TSIDParameterTuning): self.root_link_task_priority, self.root_link_task_weigth, ) + # ----------------------------- + self.controller.add_task( + self.left_foot_regularization_task, + self.left_foot_regularization_task_name, + self.left_foot_regularization_task_priority, + self.left_foot_regularization_task_weight, + ) + self.controller.add_task( + self.right_foot_regularization_task, + self.right_foot_regularization_task_name, + self.right_foot_regularization_task_priority, + self.right_foot_regularization_task_weight, + ) + + self.controller.add_task( + self.angular_momentum_task, + self.angular_momentum_task_name, + self.angular_momentum_task_priority, + self.angular_momentum_task_weight, + ) + # ----------------------------- + self.controller.finalize(self.var_handler) def run(self): @@ -398,26 +476,6 @@ def compute_com_position(self): self.com_vel = iDynTree.Twist() self.com_vel = self.kindyn.getCenterOfMassVelocity() - def update_desired_tasks( - self, - CoM_star, - CoM_dot_star, - CoM_dot_dot_star, - L, - L_Dot, - wrench_left, - wrench_right, - s_desired, - ): - self.CoM_Task.set_set_point(CoM_star, CoM_dot_star, CoM_dot_dot_star) - self.left_foot_regularization_task.set_set_point(wrench_left) - self.right_foot_regularization_task.set_set_point(wrench_right) - self.joint_regularization_task.set_set_point(s_desired) - manif_rot = blf.conversions.to_manif_rot(self.H_b[:3, :3]) - self.root_link_task.set_set_point( - manif_rot, manifpy.SO3Tangent.Zero(), manifpy.SO3Tangent.Zero() - ) - def update_task_references_mpc( self, com, @@ -428,8 +486,8 @@ def update_task_references_mpc( s_desired, wrenches_left, wrenches_right, + ang_mom_desired, ): - self.CoM_Task.set_set_point(com, dcom, ddcom) self.left_foot_tracking_task.set_set_point( left_foot_desired.transform, @@ -446,19 +504,15 @@ def update_task_references_mpc( right_contact=right_foot_desired.is_in_contact, ) self.joint_regularization_task.set_set_point(s_desired) - # self.angular_momentum_task.set_set_point(np.zeros(3), np.zeros(3)) + self.angular_momentum_task.set_set_point(ang_mom_desired, np.zeros(3)) wrench_desired_left = np.zeros(6) wrench_desired_right = np.zeros(6) - # mass = self.robot_model.get_total_mass() - # wrench_desired[2] = -(mass*9.81/2) - # wrench_desiredp[] wrench_desired_left[:3] = wrenches_left wrench_desired_right[:3] = wrenches_right - # self.left_foot_regularization_task.set_set_point(wrench_desired_left) - # self.right_foot_regularization_task.set_set_point(wrench_desired_right) + self.left_foot_regularization_task.set_set_point(wrench_desired_left) + self.right_foot_regularization_task.set_set_point(wrench_desired_right) def update_com_task(self): - angle = 0.2 * self.t CoM_des = self.COM # Calculate the x and y coordinates of the position vector using the radius and angle diff --git a/src/comodo/centroidalMPC/centroidalMPC.py b/src/comodo/centroidalMPC/centroidalMPC.py index b1eba3b..33142b3 100644 --- a/src/comodo/centroidalMPC/centroidalMPC.py +++ b/src/comodo/centroidalMPC/centroidalMPC.py @@ -1,16 +1,17 @@ -from comodo.abstractClasses.planner import Planner -from comodo.centroidalMPC.footPositionPlanner import FootPositionPlanner -import bipedal_locomotion_framework as blf from datetime import timedelta + +import bipedal_locomotion_framework as blf import idyntree.bindings as iDynTree -import numpy as np import matplotlib.pyplot as plt +import numpy as np + +from comodo.abstractClasses.planner import Planner +from comodo.centroidalMPC.footPositionPlanner import FootPositionPlanner from comodo.centroidalMPC.mpcParameterTuning import MPCParameterTuning class CentroidalMPC(Planner): def __init__(self, robot_model, step_length, frequency_ms=100): - self.dT = timedelta(milliseconds=frequency_ms) self.dT_in_seconds = frequency_ms / 1000 self.contact_planner = FootPositionPlanner( @@ -23,6 +24,7 @@ def __init__(self, robot_model, step_length, frequency_ms=100): self.gravity = iDynTree.Vector3() self.gravity.zero() self.gravity.setVal(2, -blf.math.StandardAccelerationOfGravitation) + self.total_mass = robot_model.get_total_mass() self.contact_planner.set_scaling_parameters( scaling=scaling, scalingPos=scalingPos, scalingPosY=scalingPosY ) @@ -33,16 +35,18 @@ def get_frequency_seconds(self): return self.dT_in_seconds def plan_trajectory(self): + # close the loop with centroidal dynamics directly with the sim model + com = self.kindyn.getCenterOfMassPosition().toNumPy() + dcom = self.kindyn.getCenterOfMassVelocity().toNumPy() - com = self.kindyn.getCenterOfMassPosition() - dcom = self.kindyn.getCenterOfMassVelocity() - Jcm = iDynTree.MatrixDynSize(6, 6 + self.robot_model.NDoF) - self.kindyn.getCentroidalTotalMomentumJacobian(Jcm) - nu = np.concatenate((self.w_b, self.s_dot)) - H = Jcm.toNumPy() @ nu - dcom = self.centroidal_integrator.get_solution()[1] + # this currently does not work + # total_ang_mom = self.kindyn.getCentroidalTotalMomentum().toNumPy() + # angular_mom = total_ang_mom[3:] / self.total_mass + + # fallback to the centroidal integrator angular_mom = self.centroidal_integrator.get_solution()[2] - self.centroidal_mpc.set_state(com.toNumPy(), dcom, angular_mom) + + self.centroidal_mpc.set_state(com, dcom, angular_mom) self.centroidal_mpc.set_reference_trajectory( self.com_traj, self.angular_mom_trak @@ -250,13 +254,12 @@ def initialize_centroidal_integrator(self, s, s_dot, H_b, w_b, t): self.centroidal_integrator.set_dynamical_system(self.centroidal_dynamics) self.set_state_with_base(s, s_dot, H_b, w_b, t) - com = self.kindyn.getCenterOfMassPosition() - dcom = self.kindyn.getCenterOfMassVelocity() - Jcm = iDynTree.MatrixDynSize(6, 6 + self.robot_model.NDoF) - self.kindyn.getCentroidalTotalMomentumJacobian(Jcm) - nu = np.concatenate((self.w_b, self.s_dot)) - H = Jcm.toNumPy() @ nu - self.centroidal_dynamics.set_state((com.toNumPy(), dcom.toNumPy(), H[3:])) + com = self.kindyn.getCenterOfMassPosition().toNumPy() + dcom = self.kindyn.getCenterOfMassVelocity().toNumPy() + total_ang_mom = ( + self.kindyn.getCentroidalTotalMomentum().toNumPy() / self.total_mass + ) + self.centroidal_dynamics.set_state((com, dcom, total_ang_mom[3:])) self.centroidal_integrator.set_integration_step(self.dT) def get_references(self): @@ -271,4 +274,12 @@ def get_references(self): forces_right = forces_right + item.force com = self.centroidal_integrator.get_solution()[0] dcom = self.centroidal_integrator.get_solution()[1] - return com, dcom, forces_left, forces_right + ang_mom = self.centroidal_integrator.get_solution()[2] + # ------------------------------------------------------------- + # scale with mass since the centroidal system assumes unit mass + forces_left *= self.total_mass + forces_right *= self.total_mass + ang_mom *= self.total_mass + # ------------------------------------------------------------- + + return com, dcom, forces_left, forces_right, ang_mom diff --git a/src/comodo/drakeSimulator/__init__.py b/src/comodo/drakeSimulator/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/comodo/drakeSimulator/drakeSimulator.py b/src/comodo/drakeSimulator/drakeSimulator.py new file mode 100644 index 0000000..c0f718e --- /dev/null +++ b/src/comodo/drakeSimulator/drakeSimulator.py @@ -0,0 +1,243 @@ +import logging +import math + +import numpy as np +from odio_urdf import xml_to_odio +from pydrake.common.eigen_geometry import Quaternion +from pydrake.geometry import MeshcatVisualizer, StartMeshcat +from pydrake.math import RigidTransform +from pydrake.multibody.meshcat import ContactVisualizer, ContactVisualizerParams +from pydrake.multibody.parsing import Parser +from pydrake.multibody.plant import AddMultibodyPlantSceneGraph +from pydrake.systems.analysis import Simulator, SimulatorStatus +from pydrake.systems.framework import ( + Context, + DiagramBuilder, + InputPort, + InputPortIndex, + OutputPortIndex, + PortDataType, + System, +) + +from comodo.abstractClasses.simulator import Simulator as SimulatorAbstract +from comodo.drakeSimulator.utils import DrakeURDFHelper + +# the style of wrapping the simulator is inspired by the approach by DrakeGymEnv: +# https://github.com/RobotLocomotion/drake/blob/master/bindings/pydrake/gym/_drake_gym_env.py#L177 + + +class DrakeSimulator(SimulatorAbstract): + def __init__(self) -> None: + self.duh = DrakeURDFHelper() + self.active_meshcat = False + super().__init__() + + def load_model(self, robot_model, s, xyz_rpy, kv_motors=None, Im=None): + # load the robot model construct the diagram and store the simulator + self.robot_model = robot_model + self.urdf_string = robot_model.urdf_string + + # convert the urdf string to be drake compatible + self.duh.load_urdf(urdf_string=self.urdf_string) + self.duh.remove_all_collisions() + self.duh.fix_not_in_joint_list(self.robot_model.joint_name_list) + self.duh.convert_xml_to_odio() + self.duh.add_acutation_tags() + self.urdf_string = self.duh.get_urdf_string() + + builder = DiagramBuilder() + self.time_step = 1e-3 + plant, scene_graph = AddMultibodyPlantSceneGraph( + builder, time_step=self.time_step + ) + parser = Parser(plant, scene_graph) + # TODO: Handle urdfs with mesh packages + # parser.package_map().Add("ergoCub", mesh_path) + robot_model_sim = parser.AddModels( + file_contents=self.urdf_string, + file_type="urdf", + )[0] + + # TODO: Use Im and Km to add motor params here + # ignoring those arguments for now + + # add the ground and the feet + self.duh.add_ground_with_friction(plant) + + # configure feet collisions + xMinMax = [-0.1, 0.1] + yMinMax = [-0.05, 0.05] + self.duh.add_soft_feet_collisions(plant, xMinMax=xMinMax, yMinMax=yMinMax) + + plant.Finalize() + builder.ExportInput(plant.get_actuation_input_port(), "control_input") + builder.ExportOutput(plant.get_state_output_port(), "state_output") + + # useful simulator details + self.nq = plant.num_positions() + self.nv = plant.num_velocities() + self.na = plant.num_actuators() + self.sim_joint_order = self.duh.get_sim_joint_order(plant, robot_model_sim) + # check if the joint ordering is the same + # logging.info( + # "Need joint mapping: ", not (sim_joint_order == robot_model.joint_name_list) + # ) + + # create useful variables + self.qpos = np.zeros(self.nq) + self.joint_torques = np.zeros(self.na) + + if self.visualize_robot_flag: + MeshcatVisualizer.AddToBuilder(builder, scene_graph, self.meshcat) + ContactVisualizer.AddToBuilder( + builder, + plant, + self.meshcat, + ContactVisualizerParams( + newtons_per_meter=1e3, newton_meters_per_meter=1e1 + ), + ) + + diagram = builder.Build() + self.simulator = Simulator(diagram) + + # now perform the setup + self._setup(s, xyz_rpy) + + def _setup(self, s, xyz_rpy): + # similar to DrakeGymEnv we need to expose the required ports + system = self.simulator.get_system() + self.state_output_port = system.GetOutputPort("state_output") + self.control_input_port = system.GetInputPort("control_input") + + self.context = self.simulator.get_mutable_context() + self.context.SetTime(0) + if not self.simulator.Initialize(): + logging.error("Drake simulator instance failed to initialize") + + self.simulator.get_system().SetDefaultContext(self.context) + + self.diagram = self.simulator.get_system() + plant = self.diagram.GetSubsystemByName("plant") + plant_context = self.diagram.GetMutableSubsystemContext(plant, self.context) + + self.set_base_pose_in_drake(xyz_rpy) + self.set_joint_vector_in_drake(s) + # set initial pose of the robot + plant.SetPositions(plant_context, self.qpos) + if self.visualize_robot_flag: + self.diagram.ForcedPublish(self.context) + + def reset(self): + logging.error("Drake simulator reset not implemented, recreate drake_instance") + pass + + def set_visualize_robot_flag(self, visualize_robot): + # pass meshcat to visualise the robot + self.visualize_robot_flag = visualize_robot + if self.visualize_robot_flag and not self.active_meshcat: + self.meshcat = StartMeshcat() + self.active_meshcat = True + pass + + # does this need to simulator specific? + def set_base_pose_in_drake(self, xyz_rpy): + base_xyz_quat = np.zeros(7) + # order -- quaternion+xyz + base_xyz_quat[:4] = self.RPY_to_quat(xyz_rpy[3], xyz_rpy[4], xyz_rpy[5]) + base_xyz_quat[4:] = xyz_rpy[:3] + self.qpos[:7] = base_xyz_quat + + # set initial positions of the joints + def set_joint_vector_in_drake(self, pos): + self.qpos[7:] = pos + pass + + def set_input(self, input): + # expose the acutation output port of the system to accept control + # inputs from the TSID and CentroidalMPC and step drake simulator + self.control_input_port.FixValue(self.context, input) + pass + + def step(self, n_step=1, visualize=True): + # TODO: does nothing currently + # self.control_input_port.FixValue(self.context, self.joint_torques) + self.simulator.AdvanceTo(self.context.get_time() + n_step * self.time_step) + if visualize and self.visualize_robot_flag: + self.diagram.ForcedPublish(self.context) + pass + + def step_with_motors(self, n_step, torque): + # implementation of motor level control inputs? + pass + + def get_base(self): + # this is for the state output return + # quat + pos + base_pos = self.state_output_port.Eval(self.context)[:7] + # normalise unnormalised quaternion + base_quat_wxyz = base_pos[:4] + H_b = RigidTransform( + Quaternion(base_quat_wxyz / (np.linalg.norm(base_quat_wxyz))), base_pos[4:] + ).GetAsMatrix4() + return H_b + + def get_base_velocity(self): + # this is for the state output return + base_vel = self.state_output_port.Eval(self.context)[self.nq : self.nq + 6] + # order -- linar_vel + angular_vel + return np.concatenate((base_vel[3:], base_vel[:3])) + + # TODO: Change from get_state to get_joint_state? + def get_state(self): + # this is for the state output return + robot_state = self.state_output_port.Eval(self.context) + robot_pos = robot_state[: self.nq] + robot_vel = robot_state[self.nq : self.nq + self.nv] + # return joint state to be coherent with the MuJoCo API + return ( + robot_pos[7:], + robot_vel[6:], + self.joint_torques, + ) + + def close(self): + pass + + def visualize_robot(self): + self.viewer.render() + + def get_simulation_time(self): + # get time from the simulator context + return self.context.get_time() + + def get_simulation_frequency(self): + # return the simulator timestep + return self.time_step + + # we will need similar stuff but with manifpy + def RPY_to_quat(self, roll, pitch, yaw): + cr = math.cos(roll / 2) + cp = math.cos(pitch / 2) + cy = math.cos(yaw / 2) + sr = math.sin(roll / 2) + sp = math.sin(pitch / 2) + sy = math.sin(yaw / 2) + + qw = cr * cp * cy + sr * sp * sy + qx = sr * cp * cy - cr * sp * sy + qy = cr * sp * cy + sr * cp * sy + qz = cr * cp * sy - sr * sp * cy + + # Note: The order for drake is different as compared to mujoco + return [qw, qx, qy, qz] + + def close_visualization(self): + self.meshcat.Delete() + logging.info( + "Drake uses meshcat for visualization. Close browser tab to close \ + the visualisation." + ) + self.reset() + pass diff --git a/src/comodo/drakeSimulator/utils.py b/src/comodo/drakeSimulator/utils.py new file mode 100644 index 0000000..88c22f5 --- /dev/null +++ b/src/comodo/drakeSimulator/utils.py @@ -0,0 +1,209 @@ +# some useful functions to load robot models +# into the drake simulator + +import logging +import os +import sys +import tempfile +import xml.etree.ElementTree as ET +from pathlib import Path + +import meshio +import numpy as np +from odio_urdf import * +from pydrake.geometry import ( + AddCompliantHydroelasticProperties, + AddContactMaterial, + AddRigidHydroelasticProperties, +) +from pydrake.geometry import Box as BoxDrake +from pydrake.geometry import HalfSpace, ProximityProperties +from pydrake.math import RigidTransform +from pydrake.multibody.plant import CoulombFriction +from tqdm import tqdm + + +class DrakeURDFHelper: + """A simple class to make a given URDF drake compatible""" + + def __init__(self): + pass + + def load_urdf(self, urdf_path=None, urdf_string=None, mesh_path=None): + # read the urdf + self.urdf_path = str(urdf_path) + # save the string temporarily to parse the xml + # TODO: This function blows up if not careful during an optimisation + if urdf_string is not None: + tmpfile = tempfile.NamedTemporaryFile(mode="w+") + tmpfile.write(urdf_string) + self.urdf_path = tmpfile.name + + self.mesh_path = mesh_path + self.urdf_out_path = os.getcwd() + "/urdfs" + + # load the xml from the urdf_string + tree = ET.parse(self.urdf_path) + self.root = tree.getroot() + + def convert_meshes_to_obj(self): + """Currently drake only loads .obj files and not .stl files. + This is a simple routine to allow this conversion""" + # read the meshes and convert them from .stl to .obj so that they can be loaded to Drake + logging.info("Converting all the meshes") + for child in tqdm( + self.root.findall("./link/visual/geometry/mesh") + or self.root.findall("./link/collision/geometry/mesh") + ): + # extract mesh location and name associated with the urdf + path = child.attrib["filename"] + child_mesh_path = self.mesh_path + path.split("package://")[1] + # convert the mesh and save it in the same folder + child_mesh_name = child_mesh_path.replace("stl", "obj") + if not Path(child_mesh_name).is_file(): + temp_mesh = meshio.read(child_mesh_path) + temp_mesh.write(child_mesh_name) + # replace the urdf to load .obj files + child.set("filename", path.replace("stl", "obj")) + + # # if the same, will be ignored else converted + # logging.info("Converting all the collision meshes") + # for child in tqdm(self.root.findall("./link/collision/geometry/mesh")): + # # extract mesh location and name associated with the urdf + # path = child.attrib["filename"] + # child_mesh_path = self.mesh_path + path.split("package://")[1] + # # convert the mesh and save it in the same folder + # child_mesh_name = child_mesh_path.replace("stl", "obj") + # if not Path(child_mesh_name).is_file(): + # # temp_mesh = pymesh.load_mesh(child_mesh_path) + # # pymesh.save_mesh(child_mesh_name, temp_mesh) + # temp_mesh = meshio.read(child_mesh_path) + # temp_mesh.write(child_mesh_name) + # # replace the urdf to load .obj files + # child.set("filename", path.replace("stl", "obj")) + + logging.info("Converted all the meshes from .stl to .obj") + + def remove_all_collisions(self): + """Removes all the collision tags from the URDF""" + list = [] + for link in self.root.findall("./link"): + if link.attrib["name"] not in list: + for col in link.findall("./collision"): + link.remove(col) + + def fix_not_in_joint_list(self, red_joint_name_list): + """Converts any joint not in the + reduced joint name list (red_joint_name_list) into fixed joints""" + joints = self.root.findall("./joint") + for j in joints: + if j.attrib["name"] not in red_joint_name_list: + j.attrib["type"] = "fixed" + + def convert_xml_to_odio(self): + """Converts the loaded xml to odio urdf format""" + # convert the xml to odio urdf format + self.odio_robot_dsl = xml_to_odio(self.root) + # blockPrint() + self.odio_robot = eval(self.odio_robot_dsl) + # enablePrint() + + def add_acutation_tags(self): + """Adds the actuation tags to the odio urdf""" + # extract all the non-fixed joints + joints = self.root.findall("./joint") + joint_names = [] + for j in joints: + if j.attrib["type"] != "fixed": + joint_names.append([j.attrib["name"], j.attrib["type"]]) + + # add actuation tags to the odio urdf + for j in joint_names: + actuator_name = "actuator_" + j[0] + transmission_name = "transmission" + j[0] + temp_trans = Transmission( + Type("SimpleTransmission"), + Actuator(Mechanicalreduction("1"), name=actuator_name), + Transjoint(Hardwareinterface("EffortJointInterface"), j[0]), + name=transmission_name, + ) + # Add the transmission to the robot URDF + self.odio_robot(temp_trans) + + def write_to_file(self, urdf_name): + """Writes the odio urdf to a file""" + # create a folder in the current directory called urdfs + if not os.path.exists(self.urdf_out_path): + os.makedirs(self.urdf_out_path) + filename = self.urdf_out_path + "/" + urdf_name + ".urdf" + with open(filename, "w") as f: + print(self.odio_robot, file=f) + + logging.info("Saved the urdf to {}".format(filename)) + + def get_urdf_string(self): + """Returns the urdf string""" + return str(self.odio_robot) + + def get_sim_joint_order(self, plant, robot_model_sim): + sim_joint_order = [] + for ii in plant.GetJointIndices(robot_model_sim): + jj = plant.get_joint(ii) + if jj.type_name() == "revolute": + sim_joint_order.append(jj.name()) + return sim_joint_order + + # add ground with friction + def add_ground_with_friction(self, plant): + surface_friction_ground = CoulombFriction( + static_friction=1.0, dynamic_friction=1.0 + ) + proximity_properties_ground = ProximityProperties() + AddContactMaterial( + 1e4, 1e7, surface_friction_ground, proximity_properties_ground + ) + AddRigidHydroelasticProperties(0.01, proximity_properties_ground) + + plant.RegisterCollisionGeometry( + plant.world_body(), + RigidTransform(), + HalfSpace(), + "ground_collision", + proximity_properties_ground, + ) + plant.RegisterVisualGeometry( + plant.world_body(), + RigidTransform(), + HalfSpace(), + "ground_visual", + np.array([0.5, 0.5, 0.5, 0.0]), + ) + + def add_soft_feet_collisions(self, plant, xMinMax, yMinMax): + surface_friction_feet = CoulombFriction( + static_friction=1.0, dynamic_friction=1.0 + ) + proximity_properties_feet = ProximityProperties() + AddContactMaterial(1e4, 1e7, surface_friction_feet, proximity_properties_feet) + AddCompliantHydroelasticProperties(0.01, 1e8, proximity_properties_feet) + for ii in ["l_foot_front", "l_foot_rear", "r_foot_front", "r_foot_rear"]: + # for ii in ["r_foot", "l_foot"]: + # for collision + plant.RegisterCollisionGeometry( + plant.GetBodyByName(ii), + RigidTransform(np.array([0, 0, 0])), + # RigidTransform(), + BoxDrake((xMinMax[1] - xMinMax[0]) / 2, yMinMax[1] - yMinMax[0], 2e-3), + ii + "_collision", + proximity_properties_feet, + ) + + # for visual + plant.RegisterVisualGeometry( + plant.GetBodyByName(ii), + RigidTransform(np.array([0, 0, 0])), + # RigidTransform(), + BoxDrake((xMinMax[1] - xMinMax[0]) / 2, yMinMax[1] - yMinMax[0], 2e-3), + ii + "_collision", + np.array([1.0, 1.0, 1.0, 1]), + ) diff --git a/src/comodo/robotModel/robotModel.py b/src/comodo/robotModel/robotModel.py index 98fd8d8..4b280b7 100644 --- a/src/comodo/robotModel/robotModel.py +++ b/src/comodo/robotModel/robotModel.py @@ -1,15 +1,13 @@ -from adam.casadi.computations import KinDynComputations -import numpy as np -from urchin import URDF -from urchin import Joint -from urchin import Link -import mujoco +import copy import tempfile import xml.etree.ElementTree as ET -import idyntree.bindings as iDynTree + import casadi as cs -import copy -import xml.etree.ElementTree as ET +import idyntree.bindings as iDynTree +import mujoco +import numpy as np +from adam.casadi.computations import KinDynComputations +from urchin import URDF, Joint, Link class RobotModel(KinDynComputations): @@ -80,7 +78,11 @@ def compute_desired_position_walking(self): shoulder_roll = 0.251 elbow = 0.616 - p_opts = {} + p_opts = { + "ipopt.print_level": 0, + "print_time": 0, + "ipopt.sb": "yes", + } s_opts = {"linear_solver": "mumps"} self.solver = cs.Opti() self.solver.solver("ipopt", p_opts, s_opts) @@ -179,7 +181,7 @@ def rotation_matrix_to_quaternion(self, R): z = (R[1, 0] - R[0, 1]) / S # Normalize the quaternion - length = cs.sqrt(w ** 2 + x ** 2 + y ** 2 + z ** 2) + length = cs.sqrt(w**2 + x**2 + y**2 + z**2) w /= length x /= length y /= length