Feature/rl framework integration

src/bitbots_misc/bitbots_bringup/launch/motion.launch

@@ -2,9 +2,8 @@
 <launch>
     <arg name="sim" default="false"/>
     <arg name="viz" default="false"/>
-    <arg name="walking" default="true" description="start walking" />
     <arg name="torqueless_mode" default="false" description="start without torque, for example for testing the falling detection"/>
-    <arg name="tts" default="true" description="Whether to speak" />
+    <arg name="rl_motion" default="true" description="Whether to use the RL motion system instead of the regular one" />
 
     <!-- launch the hardware interface -->
     <group unless="$(var sim)">
@@ -32,14 +31,10 @@
         <arg name="sim" value="$(var sim)" />
     </include>
 
-    <!-- launch the walking -->
-    <group if="$(var walking)">
-        <node pkg="bitbots_rl_walk" exec="walk" output="screen">
-            <param name="use_sim_time" value="$(var sim)"/>
-        </node>
-        <node pkg="bitbots_rl_walk" exec="phase_from_transform" output="screen">
-            <param name="use_sim_time" value="$(var sim)"/>
-        </node>
+    <group if="$(var rl_motion)">
+        <include file="$(find-pkg-share bitbots_rl_motion)/launch/rl_motion.launch">
+            <arg name="sim" value="$(var sim)"/>
+        </include>
     </group>
 
     <!-- launch the hcm -->

src/bitbots_misc/bitbots_bringup/launch/motion_standalone.launch

@@ -3,7 +3,7 @@
     <arg name="sim" default="false"/>
     <arg name="viz" default="false"/>
     <arg name="torqueless_mode" default="false" description="start without torque, for example for testing the falling detection"/>
-    <arg name="tts" default="true" description="Whether to enable text-to-speech"/>
+    <arg name="tts" default="false" description="Whether to enable text-to-speech"/>
 
     <include file="$(find-pkg-share bitbots_parameter_blackboard)/launch/parameter_blackboard.launch">
         <arg name="sim" value="$(var sim)"/>

src/bitbots_misc/bitbots_bringup/launch/simulator_teamplayer.launch

@@ -11,6 +11,7 @@
     <arg name="teamcom" default="false" description="Whether the team communication system should be started" />
     <arg name="vision" default="true" description="Whether the vision system should be started" />
     <arg name="world_model" default="true" description="Whether the world model should be started"/>
+    <arg name="rl_motion" default="true" description="Whether to use the RL motion system instead of the regular one" />
 
     <!-- load the general simulator -->
     <include file="$(find-pkg-share bitbots_mujoco_sim)/launch/simulator.launch" />
@@ -29,5 +30,6 @@
         <arg name="teamcom" value="$(var teamcom)" />
         <arg name="vision" value="$(var vision)"/>
         <arg name="world_model" value="$(var world_model)"/>
+        <arg name="rl_motion" value="$(var rl_motion)"/>
     </include>
 </launch>

src/bitbots_misc/bitbots_bringup/launch/teamplayer.launch

@@ -28,7 +28,7 @@
 
     <!-- load the text to speech engine -->
     <group if="$(var tts)">
-        <include file="$(find-pkg-share bitbots_tts)/launch/tts.launch"/>
+        <include file="$(find-pkg-share bitbots_tts)/launch/tts.launch" />
     </group >
 
     <!-- load the diagnostic aggregator -->
@@ -42,8 +42,7 @@
     <!-- load the motion -->
     <group if="$(var motion)">
         <include file="$(find-pkg-share bitbots_bringup)/launch/motion.launch">
-            <arg name="sim" value="$(var sim)"/>
-            <arg name="tts" value="$(var tts)" />
+            <arg name="sim" value="$(var sim)" />
         </include>
     </group>
 

src/bitbots_motion/bitbots_rl_motion/bitbots_rl_motion/phase.py

@@ -1,26 +1,15 @@
-from typing import Optional
-
 import numpy as np
+from rclpy.node import Node
 
 
 class Phase:
-    _phase: np.ndarray = np.array([0.0, np.pi], dtype=np.float32)
-    _phase_dt: Optional[float]
-
-    def __init__(self, node):
+    def __init__(self, node: Node):
         self._node = node
-
-        if self._node.get_parameter("phase.use_phase").value:
-            self._control_dt = self._node.get_parameter("phase.control_dt").value
-            self._gait_frequency = self._node.get_parameter("phase.gait_frequency").value
-            self._phase_dt = 2 * np.pi * self._gait_frequency * self._control_dt
-            self._use_phase = True
-        else:
-            self._control_dt = None
-            self._gait_frequency = None
-            self._phase_dt = None
-            self._use_phase = False
-            self._node.get_logger().warning("No phase was found! Using policy without phase!")
+        self._use_phase = self._node.get_parameter("phase.use_phase").value
+        self._phase = np.array(self._node.get_parameter("phase.initial_phase").value)
+        self._control_dt = self._node.get_parameter("phase.control_dt").value
+        self._gait_frequency = self._node.get_parameter("phase.gait_frequency").value
+        self._phase_dt = 2 * np.pi * self._gait_frequency * self._control_dt
 
         self._obs_phase = None
 

src/bitbots_motion/bitbots_rl_motion/configs/mjlab_walk_model.yaml

@@ -0,0 +1,87 @@
+mjlab_walk_node:
+  ros__parameters:
+    model: walk_mjlab.onnx
+
+    joints:
+      ordered_relevant_joint_names: [
+        "r_hip_pitch_joint",
+        "l_hip_pitch_joint",
+        "r_hip_roll_joint",
+        "l_hip_roll_joint",
+        "r_thigh_joint",
+        "l_thigh_joint",
+        "r_calf_joint",
+        "l_calf_joint",
+        "r_ankle_pitch_joint",
+        "l_ankle_pitch_joint",
+        "r_ankle_roll_joint",
+        "l_ankle_roll_joint",
+        "r_shoulder_roll_joint",
+        "l_shoulder_roll_joint",
+        "r_shoulder_pitch_joint",
+        "l_shoulder_pitch_joint",
+        "r_upper_arm_joint",
+        "l_upper_arm_joint",
+        "r_elbow_joint",
+        "l_elbow_joint",
+      ] 
+      walkready_state: [
+        0.25, -0.25,
+        0.0, 0.0,
+        0.0, 0.0,
+        0.65, -0.65,
+        0.4, -0.4,
+        0.0, 0.0,
+        0.0, 0.0,
+        0.0, 0.0,
+        0.0, 0.0,
+        0.0, 0.0,
+      ]
+      kp: [
+        # legs
+        55.0, 55.0, 
+        55.0, 55.0, 
+        55.0, 55.0, 
+        55.0, 55.0, 
+        55.0, 55.0, 
+        55.0, 55.0, 
+        # arms
+        6.0, 6.0,
+        6.0, 6.0,
+        6.0, 6.0,
+        6.0, 6.0,
+      ]
+      kd: [
+        # legs
+        1.1, 1.1, 
+        1.4, 1.4,
+        1.1, 1.1,
+        1.1, 1.1,
+        1.1, 1.1,
+        1.1, 1.1,
+        # arms
+        0.6, 0.6,
+        0.6, 0.6,
+        0.6, 0.6,
+        0.6, 0.6,
+      ]
+      action_scales: [
+        # legs
+        0.5, 0.5,
+        0.5, 0.5,
+        0.5, 0.5,
+        0.5, 0.5,
+        0.5, 0.5,
+        0.5, 0.5,
+        # arms
+        0.5, 0.5,
+        0.5, 0.5,
+        0.5, 0.5,
+        0.5, 0.5
+      ]
+    phase:
+      control_dt: 0.02
+      gait_frequency: 99.99
+      use_phase: False
+
+    providers: ["CPUExecutionProvider"]

src/bitbots_motion/bitbots_rl_motion/configs/playground_walk_model.yaml

@@ -0,0 +1,57 @@
+walk_node:
+  ros__parameters:
+    model: policy_walk.onnx
+
+    joints:
+      ordered_relevant_joint_names: [
+        "r_hip_pitch_joint",
+        "r_hip_roll_joint",
+        "r_thigh_joint",
+        "r_calf_joint",
+        "r_ankle_pitch_joint",
+        "r_ankle_roll_joint",
+        "l_hip_pitch_joint",
+        "l_hip_roll_joint",
+        "l_thigh_joint",
+        "l_calf_joint",
+        "l_ankle_pitch_joint",
+        "l_ankle_roll_joint"
+      ] 
+      walkready_state: [
+         0.6, 0.0, 0.0,  1.2,  0.6, 0.0,
+        -0.6, 0.0, 0.0, -1.2, -0.6, 0.0
+      ]
+      kp: [
+        # legs
+        30.0, 30.0, 
+        30.0, 30.0, 
+        30.0, 30.0, 
+        30.0, 30.0, 
+        30.0, 30.0, 
+        30.0, 30.0
+      ]
+      kd: [
+        # legs
+        1.1, 1.1, 
+        1.1, 1.1,
+        1.1, 1.1,
+        1.1, 1.1,
+        1.1, 1.1,
+        1.1, 1.1
+      ]
+      action_scales: [
+        # legs
+        0.5, 0.5,
+        0.5, 0.5,
+        0.5, 0.5,
+        0.5, 0.5,
+        0.5, 0.5,
+        0.5, 0.5
+      ]
+    phase:
+      use_phase: True
+      control_dt: 0.02
+      gait_frequency: 0.75
+      initial_phase: [1.5707963267948966, -1.5707963267948966]
+
+    providers: ["CPUExecutionProvider"]

src/bitbots_motion/bitbots_rl_motion/handlers/command_handler.py

@@ -9,15 +9,27 @@
 class CommandHandler(Handler):
     def __init__(self, node):
         self._node = node
-
         self._cmd_vel: Optional[Twist] = None
-
         self._cmd_vel_sub = self._node.create_subscription(Twist, "cmd_vel", self._cmd_vel_callback, 10)
 
     def get_command(self):
         assert self._cmd_vel is not None
-        command = np.array([self._cmd_vel.linear.x, self._cmd_vel.linear.y, self._cmd_vel.angular.z], dtype=np.float32)
-        return command
+        return np.array(
+            [self._cmd_vel.linear.x, self._cmd_vel.linear.y, self._cmd_vel.angular.z],
+            dtype=np.float32,
+        )
+
+    def get_command_with_stop(self):
+        assert self._cmd_vel is not None
+        stop_signal = float(self.get_stop_signal())
+        return np.array(
+            [self._cmd_vel.linear.x, self._cmd_vel.linear.y, self._cmd_vel.angular.z, stop_signal],
+            dtype=np.float32,
+        )
+
+    def get_stop_signal(self) -> bool:
+        assert self._cmd_vel is not None
+        return self._cmd_vel.angular.x != 0.0
 
     def has_data(self):
         return self._cmd_vel is not None

src/bitbots_motion/bitbots_rl_motion/handlers/gravity_handler.py

@@ -2,7 +2,6 @@
 
 import numpy as np
 from sensor_msgs.msg import Imu
-from transforms3d.euler import euler2mat
 from transforms3d.quaternions import quat2mat
 
 from handlers.handler import Handler
@@ -17,23 +16,17 @@ def __init__(self, node):
         self._imu_sub = self._node.create_subscription(Imu, "imu/data", self._imu_callback, 10)
 
     # Callables
-    def _imu_callback(self, msg):
+    def _imu_callback(self, msg: Imu):
         self._imu_data = msg
 
-    def has_data(self):
+    def has_data(self) -> bool:
         return self._imu_data is not None
 
-    def get_gravity(self):
-        assert self._imu_data is not None
-        gravity = (
-            quat2mat(
-                [
-                    self._imu_data.orientation.w,
-                    self._imu_data.orientation.x,
-                    self._imu_data.orientation.y,
-                    self._imu_data.orientation.z,
-                ]
-            )
-            @ euler2mat(0, -0.0, 0)
-        ).T @ np.array([0, 0, -1], dtype=np.float32)
+    def get_gravity(self) -> np.ndarray:
+        """
+        Returns the gravity vector in the robot's base frame computed from the IMU orientation.
+        """
+        assert self.has_data(), "IMU data is not available yet"
+        q = self._imu_data.orientation
+        gravity = quat2mat([q.w, q.x, q.y, q.z]).T @ np.array([0, 0, -1], dtype=np.float32)
         return gravity