implementing rl kick in behaviour

Author: sWintermoor

src/bitbots_behavior/bitbots_blackboard/bitbots_blackboard/capsules/kick_capsule.py

@@ -1,6 +1,7 @@
 from enum import Flag
 from typing import Optional
 
+from geometry_msgs.msg import PoseStamped
 from rclpy.action import ActionClient
 from rclpy.callback_groups import ReentrantCallbackGroup
 from rclpy.duration import Duration
@@ -34,6 +35,7 @@ class WalkKickTargets(Flag):
         RIGHT = True
 
     walk_kick_pub: Publisher
+    rl_kick_pub: Publisher
 
     def __init__(self, node, blackboard):
         super().__init__(node, blackboard)
@@ -42,6 +44,7 @@ def __init__(self, node, blackboard):
         """
         self.walk_kick_pub = self._node.create_publisher(Bool, "/kick", 1)
         # self.connect_dynamic_kick()  Do not connect if dynamic_kick is disabled
+        self.rl_kick_pub = self._node.create_publisher(PoseStamped, "/rl_kick/goal", 1)
 
     def walk_kick(self, target: WalkKickTargets) -> None:
         """
@@ -81,6 +84,13 @@ def dynamic_kick(self, goal: Kick.Goal) -> None:
         self.last_goal = goal
         self.last_goal_sent = self._node.get_clock().now()
 
+    def rl_kick(self, goal_pose: PoseStamped) -> None:
+        """
+        Kick the ball using the RL kick
+        :param target: Target for the RL kick (e.g. left or right foot)
+        """
+        self.rl_kick_pub.publish(goal_pose)
+
     def __feedback_cb(self, feedback):
         self.last_feedback: Kick.Feedback = feedback.feedback
         self.last_feedback_received = self._node.get_clock().now()

src/bitbots_behavior/bitbots_body_behavior/bitbots_body_behavior/behavior_dsd/actions/kick_ball.py

@@ -123,3 +123,39 @@ def perform(self, reevaluate=False):
                 self._goal_sent = True
             else:
                 self.pop()
+
+
+class RLKick(AbstractKickAction):
+    def __init__(self, blackboard, dsd, parameters):
+        super().__init__(blackboard, dsd, parameters)
+
+        self.kick_length = self.blackboard.config["kick_cost_kick_length"]
+        self.angular_range = self.blackboard.config["kick_cost_angular_range"]
+        self.max_kick_angle = self.blackboard.config["max_kick_angle"]
+        self.num_kick_angles = self.blackboard.config["num_kick_angles"]
+        self.penalty_kick_angle = self.blackboard.config["penalty_kick_angle"]
+
+    def perform(self, reevaluate=False):
+        goal = Kick.Goal()
+        goal.header.stamp = self.blackboard.node.get_clock().now().to_msg()
+        goal.header.frame_id = self.blackboard.world_model.base_footprint_frame
+
+        ball_u, ball_v = self.blackboard.world_model.get_ball_position_uv()
+        goal.kick_speed = 1.0
+        goal.ball_position.x = ball_u
+        goal.ball_position.y = ball_v
+        goal.ball_position.z = 0.0
+        goal.unstable = False
+
+        kick_direction = self.blackboard.costmap.get_best_kick_direction(
+            -self.max_kick_angle,
+            self.max_kick_angle,
+            self.num_kick_angles,
+            self.kick_length,
+            self.angular_range,
+        )
+
+        goal.kick_direction = quat_from_yaw(kick_direction)
+
+        self.blackboard.kick.rl_kick(goal)
+        self.pop()

…ots_rl_walk/bitbots_rl_walk/real_walk.py …/bitbots_rl_walk/bitbots_rl_walk/kick.pysrc/bitbots_rl_walk/bitbots_rl_walk/real_walk.py renamed to src/bitbots_rl_walk/bitbots_rl_walk/kick.py src/bitbots_rl_walk/bitbots_rl_walk/real_walk.py renamed to src/bitbots_rl_walk/bitbots_rl_walk/kick.py

@@ -21,7 +21,7 @@
 import numpy as np
 import onnxruntime as rt
 from ament_index_python import get_package_share_directory
-from geometry_msgs.msg import Twist
+from geometry_msgs.msg import PoseStamped, Twist
 from rclpy.node import Node
 from sensor_msgs.msg import Imu, JointState
 from transforms3d.euler import euler2mat
@@ -153,18 +153,19 @@
 ]
 
 
-class WalkNode(Node):
+class KickNode(Node):
     """Node to control the wolfgang humanoid."""
 
     _previous_action: np.ndarray = np.zeros(len(ORDERED_RELEVANT_JOINT_NAMES), dtype=np.float32)
     _imu_data: Optional[Imu] = None
     _joint_state: Optional[JointState] = None
     _cmd_vel: Optional[Twist] = None
+    _goal_pose: Optional[PoseStamped] = None
     _phase: np.ndarray = np.array([0.0, np.pi], dtype=np.float32)
     _phase_dt: float
 
     def __init__(self):
-        super().__init__("reinforcement_learning_walk_inference_node")
+        super().__init__("reinforcement_learning_kick_inference_node")
 
         # Set sim time parameter to true
         # self.set_parameters([
@@ -179,6 +180,7 @@ def __init__(self):
         self._imu_sub = self.create_subscription(Imu, "imu/data", self._imu_callback, 10)
         self._joint_state_sub = self.create_subscription(JointState, "joint_states", self._joint_state_callback, 10)
         self._cmd_vel_sub = self.create_subscription(Twist, "cmd_vel", self._cmd_vel_callback, 10)
+        self._goal_pose_sub = self.create_subscription(PoseStamped, "goal_pose", self._goal_pose_callback, 10)
 
         self._timer = self.create_timer(CONTROL_DT, self._timer_callback)
 
@@ -211,6 +213,9 @@ def _joint_state_callback(self, msg: JointState):
     def _cmd_vel_callback(self, msg: Twist):
         self._cmd_vel = msg
 
+    def _goal_pose_callback(self, msg: PoseStamped):
+        self._goal_pose = msg
+
     def _imu_callback(self, msg: Imu):
         self._imu_data = msg
 
@@ -274,6 +279,18 @@ def _timer_callback(self):
 
         command = np.array([self._cmd_vel.linear.x, self._cmd_vel.linear.y, self._cmd_vel.angular.z], dtype=np.float32)
 
+        """
+        rel_ball_pos = np.array(
+            [
+                self._ball_pos_rel_filter.pose.pose.position.x,
+                self._ball_pos_rel_filter.pose.pose.position.y,
+            ],
+            dtype=np.float32,
+        )
+
+        rel_target_pos = np.array([])
+        """
+
         obs = np.hstack(
             [
                 gyro,  # 3
@@ -283,6 +300,7 @@ def _timer_callback(self):
                 joint_velocities,  # 18
                 self._previous_action,  # 18  # Previous action
                 phase,  # 2
+                # rel_ball_pos,  # 2
             ]
         ).astype(np.float32)
 
@@ -310,7 +328,7 @@ def main():
     import rclpy
 
     rclpy.init()
-    node = WalkNode()
+    node = KickNode()
     rclpy.spin(node)
     node.destroy_node()
     rclpy.try_shutdown()

src/bitbots_rl_walk/bitbots_rl_walk/walk.py

@@ -21,16 +21,15 @@
 import numpy as np
 import onnxruntime as rt
 from ament_index_python import get_package_share_directory
-from geometry_msgs.msg import PoseWithCovarianceStamped, Twist
+from geometry_msgs.msg import Twist
 from rclpy.node import Node
 from sensor_msgs.msg import Imu, JointState
-from soccer_vision_3d_msgs.msg import GoalpostArray
 from transforms3d.euler import euler2mat
 from transforms3d.quaternions import quat2mat
 
 from bitbots_msgs.msg import JointCommand
 
-ONNX_MODEL = os.path.join(get_package_share_directory("bitbots_rl_walk"), "models", "wolfgang_walk_ppo.onnx")
+ONNX_MODEL = os.path.join(get_package_share_directory("bitbots_rl_walk"), "models", "wolfgang_kick_ppo.onnx")
 
 PREPARATION_STATE = np.array(
     [
@@ -161,8 +160,6 @@ class WalkNode(Node):
     _imu_data: Optional[Imu] = None
     _joint_state: Optional[JointState] = None
     _cmd_vel: Optional[Twist] = None
-    _ball_pos_rel_filter: Optional[PoseWithCovarianceStamped] = None
-    _goal_posts_rel: Optional[GoalpostArray] = None
     _phase: np.ndarray = np.array([0.0, np.pi], dtype=np.float32)
     _phase_dt: float
 
@@ -182,12 +179,6 @@ def __init__(self):
         self._imu_sub = self.create_subscription(Imu, "imu/data", self._imu_callback, 10)
         self._joint_state_sub = self.create_subscription(JointState, "joint_states", self._joint_state_callback, 10)
         self._cmd_vel_sub = self.create_subscription(Twist, "cmd_vel", self._cmd_vel_callback, 10)
-        self._ball_pos_rel_filter_sub = self.create_subscription(
-            PoseWithCovarianceStamped, "ball_position_relative_filtered", self._ball_pos_rel_filter_callback, 10
-        )
-        self._goal_posts_rel_sub = self.create_subscription(
-            GoalpostArray, "goal_posts_relative", self._goal_posts_rel_callback, 10
-        )
 
         self._timer = self.create_timer(CONTROL_DT, self._timer_callback)
 
@@ -223,12 +214,6 @@ def _cmd_vel_callback(self, msg: Twist):
     def _imu_callback(self, msg: Imu):
         self._imu_data = msg
 
-    def _ball_pos_rel_filter_callback(self, msg: PoseWithCovarianceStamped):
-        self._ball_pos_rel_filter = msg
-
-    def _goal_posts_rel_callback(self, msg: GoalpostArray):
-        self._goal_posts_rel = msg
-
     def _timer_callback(self):
         """Timer callback to publish joint commands based on the ONNX policy."""
         if self._imu_data is None or self._joint_state is None or self._cmd_vel is None:
@@ -289,18 +274,6 @@ def _timer_callback(self):
 
         command = np.array([self._cmd_vel.linear.x, self._cmd_vel.linear.y, self._cmd_vel.angular.z], dtype=np.float32)
 
-        """
-        rel_ball_pos = np.array(
-            [
-                self._ball_pos_rel_filter.pose.pose.position.x,
-                self._ball_pos_rel_filter.pose.pose.position.y,
-            ],
-            dtype=np.float32,
-        )
-
-        rel_target_pos = np.array([])
-        """
-
         obs = np.hstack(
             [
                 gyro,  # 3
@@ -310,7 +283,6 @@ def _timer_callback(self):
                 joint_velocities,  # 18
                 self._previous_action,  # 18  # Previous action
                 phase,  # 2
-                # rel_ball_pos,  # 2
             ]
         ).astype(np.float32)
 

src/bitbots_rl_walk/setup.py

@@ -9,7 +9,10 @@
     data_files=[
         ("share/ament_index/resource_index/packages", ["resource/" + package_name]),
         ("share/" + package_name, ["package.xml"]),
-        ("share/" + package_name + "/models", ["models/wolfgang_kick_ppo.onnx"]),
+        (
+            "share/" + package_name + "/models",
+            ["models/wolfgang_policy.onnx", "models/wolfgang_walk_ppo.onnx", "models/wolfgang_kick_ppo.onnx"],
+        ),
     ],
     install_requires=["setuptools"],
     zip_safe=True,