kick node

Author: sWintermoor

src/bitbots_motion/bitbots_rl_motion/bitbots_rl_motion/phase.py

@@ -14,6 +14,9 @@ def __init__(self, config):
         self._gait_frequency = self._config["phase"]["gait_frequency"]
         self._phase_dt = 2 * np.pi * self._gait_frequency * self._control_dt
 
+    def set_phase(self, new_phase):
+        self._phase = new_phase
+
     def get_phase(self):
         return self._phase
 

src/bitbots_motion/bitbots_rl_motion/bitbots_rl_motion/policy_nodes.py

@@ -2,19 +2,18 @@
 
 import rclpy
 from ament_index_python import get_package_share_directory
-from nodes.walk_node import WalkNode
+from nodes.kick_node import KickNode
 
 
 def main():
     rclpy.init()
 
     wolfgang_config = os.path.join(get_package_share_directory("bitbots_rl_motion"), "config", "wolfgang_config.yaml")
-    # kick_policy_path = os.path.join(get_package_share_directory("bitbots_rl_motion"), "models", "wolfgang_kick_ppo.onnx")
 
-    walk_node = WalkNode(wolfgang_config)
-    # kick_node = RLNode(kick_policy_path)
+    # walk_node = WalkNode(wolfgang_config)
+    kick_node = KickNode(wolfgang_config)
 
-    rclpy.spin(walk_node)
-    walk_node.destroy()
+    rclpy.spin(kick_node)
+    kick_node.destroy()
 
     rclpy.try_shutdown()

src/bitbots_motion/bitbots_rl_motion/configs/wolfgang_config.yaml

@@ -1,5 +1,6 @@
 models: 
   walk_model: wolfang_walk_ppo.onnx
+  kick_model: wolfgang_kick_ppo.onnx
 
 joints:
   ordered_relevant_joint_names: [

src/bitbots_motion/bitbots_rl_motion/handlers/ball_handler.py

@@ -0,0 +1,13 @@
+from bitbots_rl_motion.handlers.handler import Handler
+
+
+class BallHandler(Handler):
+    def __init__(self, config):
+        super().__init__(config)
+        self.ball_pos = None
+
+    def ball_pos_callback(self, msg):
+        self.ball_pos = msg
+
+    def get_ball_pos(self):
+        return self.ball_pos

src/bitbots_motion/bitbots_rl_motion/nodes/kick_node.py

@@ -0,0 +1,70 @@
+import numpy as np
+from bitbots_rl_motion.handlers.ball_handler import BallHandler
+from bitbots_rl_motion.handlers.joint_handler import JointHandler
+from bitbots_rl_motion.nodes.rl_node import RLNode
+from geometry_msgs.msg import PoseStamped
+from handlers.gravity_handler import GravityHandler
+from handlers.gyro_handler import GyroHandler
+from sensor_msgs.msg import Imu, JointState
+
+from bitbots_msgs.msg import JointCommand
+
+
+class KickNode(RLNode):
+    def __init__(self, config_path: str):
+        super().__init__(config_path)
+
+        # loading model
+        model = self._config["models"]["kick_model"]
+        self.load_model(model)
+
+        # publishers
+        self._joint_command_pub = self.create_publisher(JointCommand, "walking_motor_goals", 10)
+
+        # subscribers
+        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._ball_pos_sub = self.create_subscription(PoseStamped, "ball_pos", self._ball_pos_callback, 10)
+
+        # handlers
+        self._gyro_handler = GyroHandler(self._config)
+        self._gravity_handler = GravityHandler(self._config)
+        self._joint_handler = JointHandler(self._config)
+        self._ball_handler = BallHandler(self._config)
+
+        # observations
+
+        self._obs = np.hstack(
+            [
+                self._gyro_handler.get_gyro(),
+                self._gravity_handler.get_gravity(),
+                self._joint_handler.get_velocity_data(),
+                self._joint_handler.get_angle_data(),
+                self._joint_handler.get_previous_action(),
+                self._phase.get_phase(),
+                self._ball_handler.get_ball_pos(),
+            ]
+        ).astype(np.float32)
+
+    # callback functions
+
+    def _imu_callback(self, msg):
+        self._gyro_handler.imu_callback(msg)
+        self._gravity_handler.imu_callback(msg)
+
+    def _joint_state_callback(self, msg):
+        self._joint_handler.joint_state_callback(msg)
+
+    def _ball_pos_callback(self, msg):
+        self._ball_handler.ball_pos_callback(msg)
+
+    # load phase function
+
+    def load_phase(self):
+        pass
+
+    # publisher function
+
+    def publisher(self, onnx_pred):
+        joint_command = self._joint_handler.get_joint_commands(onnx_pred)
+        self._joint_command_pub.publish(joint_command)

src/bitbots_motion/bitbots_rl_motion/nodes/rl_node.py

@@ -23,6 +23,7 @@
 import onnxruntime as rt
 import yaml
 from ament_index_python import get_package_share_directory
+from bitbots_rl_motion.bitbots_rl_motion.phase import PhaseObject
 from rclpy.node import Node
 from rclpy.qos import QoSProfile
 from rclpy.subscription import Subscription
@@ -44,8 +45,8 @@ class SubscriptionParam(NamedTuple):
 
     def __init__(self, config_path: str):
         self._config = self._load_config(config_path)
+        self._phase = PhaseObject(self._config)
         self._obs = None  # should be defined in subclass
-        self._phase_handler = None  # should be defined in subclass
 
         self._timer = self.create_timer(self._config["phase"]["control_dt"], self._timer_callback)
         self.load_phase()
@@ -77,9 +78,9 @@ def _timer_callback(self):
 
             # TODO consider IMU mounting offset
 
-            self._phase_handler.set_obs_phase(
+            self._phase.set_phase(
                 np.array(
-                    [np.cos(self._phase_handler.get_phase()), np.sin(self._phase_handler.get_phase())],
+                    [np.cos(self._phase.get_phase()), np.sin(self._phase.get_phase())],
                     dtype=np.float32,
                 ).flatten()
             )
@@ -91,8 +92,8 @@ def _timer_callback(self):
 
             self.publisher(onnx_pred)
 
-            phase_tp1 = self._phase_handler.get_phase() + self._phase_handler.get_phase_dt()
-            self._phase_handler.set_phase(np.fmod(phase_tp1 + np.pi, 2 * np.pi) - np.pi)
+            phase_tp1 = self._phase.get_phase() + self._phase.get_phase_dt()
+            self._phase.set_phase(np.fmod(phase_tp1 + np.pi, 2 * np.pi) - np.pi)
         else:
             raise ConfigError("Configuration is missing! Try to run self.config() in init.")
 

src/bitbots_motion/bitbots_rl_motion/nodes/walk_node.py

@@ -2,7 +2,6 @@
 
 import numpy as np
 from bitbots_rl_motion.handlers.joint_handler import JointHandler
-from bitbots_rl_motion.handlers.phase_handler import PhaseHandler
 from bitbots_rl_motion.nodes.rl_node import RLNode
 from geometry_msgs.msg import Twist
 from handlers.command_handler import CommandHandler
@@ -34,7 +33,6 @@ def __init__(self, config_path: str):
         self._gravity_handler = GravityHandler(self._config)
         self._joint_handler = JointHandler(self._config)
         self._command_handler = CommandHandler(self._config)
-        self._phase_handler = PhaseHandler(self._config)
 
         # observations
 
@@ -46,7 +44,7 @@ def __init__(self, config_path: str):
                 self._joint_handler.get_velocity_data(),
                 self._joint_handler.get_angle_data(),
                 self._joint_handler.get_previous_action(),
-                self._phase_handler.get_phase(),  # TODO: Check whether correct
+                self._phase.get_phase(),
             ]
         ).astype(np.float32)
 

src/bitbots_motion/bitbots_rl_motion/rl_motion.launch

@@ -0,0 +1,2 @@
+<launch>
+    

src/bitbots_motion/bitbots_rl_motion/setup.py

@@ -18,17 +18,14 @@
     ],
     install_requires=["setuptools"],
     zip_safe=True,
-    maintainer="florian",
-    maintainer_email="git@flova.de",
+    maintainer="mark oliver",
+    maintainer_email="git@sWintermoor.de",
     description="TODO: Package description",
     license="TODO: License declaration",
     tests_require=["pytest"],
     entry_points={
         "console_scripts": [
-            "walk = bitbots_rl_motion.walk:main",
-            "kick = bitbots_rl_motion.kick:main",
-            "walk_kick = bitbots_rl_motion.walk_kick:main",
-            "forward_kick = bitbots_rl_motion.forward_kick:main",
+            "run_policies = bitbots_rl_motion.policy_nodes:main",
         ],
     },
 )