phase in child class

Author: sWintermoor

src/bitbots_rl_walk/confg/timer_phase_confg.py

@@ -0,0 +1,27 @@
+import numpy as np
+
+CONTROL_DT = 0.02  # Control loop frequency in seconds
+GAIT_FREQUENCY = 1.5  # Gait frequency in Hz
+
+
+class TimerPhaseConfg:
+    _phase: np.ndarray = np.array([0.0, np.pi], dtype=np.float32)
+    _phase_dt: float
+
+    def __init__(self):
+        # Phase time
+        self._control_dt = CONTROL_DT
+        self._gait_frequency = GAIT_FREQUENCY
+        self._phase_dt = 2 * np.pi * GAIT_FREQUENCY * CONTROL_DT
+
+    def get_control_dt(self):
+        return self._control_dt
+
+    def get_gait_frequency(self):
+        return self._gait_frequency
+
+    def get_phase(self):
+        return self._phase
+
+    def get_phase_dt(self):
+        return self._phase_dt

src/bitbots_rl_walk/handler/joint_handler.py

@@ -53,7 +53,16 @@ class JointHandler(Handler):
     def __init__(self, ordered_relevant_joint_names=ORDERED_RELEVANT_JOINT_NAMES, walkready_state=WALKREADY_STATE):
         self._ordered_relevant_joint_names = ordered_relevant_joint_names
         self._walkready_state = walkready_state
+        self._previous_action: np.ndarray = np.zeros(len(self._ordered_relevant_joint_names), dtype=np.float32)
         self._joint_state = None
+        self._obs_phase = None
+        self._phase = None
+
+    def set_obs_phase(self, phase):
+        self._obs_phase = phase
+
+    def set_phase(self, phase):
+        self._phase = phase
 
     def get_angle_data(self):
         joint_angles = (
@@ -92,6 +101,8 @@ def get_walkready_joint_command(self, timestamp):
         joint_command.header.stamp = timestamp.to_msg()
         joint_command.positions = self._walkready_state
 
+        self._previous_action = joint_command
+
         return joint_command
 
     def get_joint_commands(self, onnx_pred):
@@ -103,7 +114,18 @@ def get_joint_commands(self, onnx_pred):
         joint_command.accelerations = [-1.0] * len(self._ordered_relevant_joint_names)
         joint_command.max_currents = [-1.0] * len(self._ordered_relevant_joint_names)
 
+        self._previous_action = joint_command
+
         return joint_command
 
+    def get_previous_action(self):
+        return self._previous_action
+
+    def get_obs_phase(self):
+        return self._obs_phase
+
+    def get_phase(self):
+        return self._phase
+
     def joint_state_callback(self, msg):
         self._joint_state = msg

src/bitbots_rl_walk/nodes/rl_node.py

@@ -31,10 +31,6 @@
 class RLNode(Node):
     """Node to control the wolfgang humanoid."""
 
-    # TODO: _previous_action: np.ndarray = np.zeros(len(ORDERED_RELEVANT_JOINT_NAMES), dtype=np.float32)
-    _phase: np.ndarray = np.array([0.0, np.pi], dtype=np.float32)
-    _phase_dt: float
-
     class PublisherParam(NamedTuple):
         msg_type: int
         topic: str
@@ -63,9 +59,7 @@ def __init__(self, path_to_model):
         for out in self._onnx_model.graph.output:
             self._onnx_output_name.append(out)
 
-        # Phase time
-        self._phase_dt = 2 * np.pi * GAIT_FREQUENCY * CONTROL_DT
-
+        # TODO: Move timer to child class
         self._timer = self.create_timer(CONTROL_DT, self._timer_callback)
 
         self.load_phase()
@@ -76,6 +70,9 @@ def __init__(self, path_to_model):
             if type(value) is Subscription:
                 self._subs.append(key)
 
+        self._obs = None  # should be defined in subclass
+        self._timer_phase_confg = None  # Should be defined in subclass
+
     # TODO: fix
     def _timer_callback(self):
         for subscription in self._subs:
@@ -90,17 +87,22 @@ def _timer_callback(self):
 
         # TODO consider IMU mounting offset
 
-        self._obs_phase = np.array([np.cos(self._phase), np.sin(self._phase)], dtype=np.float32).flatten()
+        self._timer_phase_confg.set_obs_phase(
+            np.array(
+                [np.cos(self._timer_phase_confg.get_phase()), np.sin(self._timer_phase_confg.get_phase())],
+                dtype=np.float32,
+            ).flatten()
+        )
 
         # Run the ONNX model
-        onnx_input = {self._onnx_input_name[0]: self.obs().reshape(1, -1)}  # TODO: Improve input
+        onnx_input = {self._onnx_input_name[0]: self._obs.reshape(1, -1)}  # TODO: Improve input
         onnx_pred = self._onnx_session.run(self._onnx_output_name, onnx_input)[0][0]
         self._previous_action = onnx_pred
 
         self.publisher(onnx_pred)
 
-        phase_tp1 = self._phase + self._phase_dt
-        self._phase = np.fmod(phase_tp1 + np.pi, 2 * np.pi) - np.pi
+        phase_tp1 = self._timer_phase_confg.get_phase() + self._timer_phase_confg.get_phase_dt()
+        self._timer_phase_confg.set_phase(np.fmod(phase_tp1 + np.pi, 2 * np.pi) - np.pi)
 
     def obs():
         # Should be defined in subclass

src/bitbots_rl_walk/nodes/walk_node.py

@@ -8,6 +8,7 @@
 from sensor_msgs.msg import Imu, JointState
 
 from bitbots_msgs.msg import JointCommand
+from bitbots_rl_walk.confg.timer_phase_confg import TimerPhaseConfg
 from bitbots_rl_walk.handler.joint_handler import JointHandler
 from bitbots_rl_walk.nodes.rl_node import RLNode
 
@@ -28,33 +29,34 @@ def __init__(self, walk_policy_path):
         self._joint_handler = JointHandler()
         self._command_handler = CommandHandler()
 
-        super().__init__(walk_policy_path)
+        self._timer_phase_confg = TimerPhaseConfg()
 
-    def _imu_callback(self, msg):
-        self._gyro_handler.imu_callback(msg)
-        self._gravity_handler.imu_callback(msg)
+        # TODO: timer is missing
 
-    def _joint_state_callback(self, msg):
-        self._joint_handler.joint_state_callback(msg)
-
-    def _cmd_vel_callback(self, msg):
-        self._command_handler.cmd_vel_callback(msg)
-
-    def obs(self):
-        obs = np.hstack(
+        self._obs = np.hstack(
             [
                 self._gyro_handler.get_data(),  # 3
                 self._gravity_handler.get_data(),  # 4
                 self._command_handler.get_data(),  # 3
                 self._joint_handler.get_velocity_data(),  # 18
                 self._joint_handler.get_angle_data(),  # 18
                 # TODO: fix
-                self._previous_action,  # 18  # Previous action
-                self._obs_phase,  # 2
+                self._joint_handler.get_previous_action(),  # 18  # Previous action
+                self._timer_phase_confg.get_obs_phase(),  # 2
             ]
         ).astype(np.float32)
 
-        return obs
+        super().__init__(walk_policy_path)
+
+    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 _cmd_vel_callback(self, msg):
+        self._command_handler.cmd_vel_callback(msg)
 
     def load_phase(self):
         walkready_command = self._joint_handler.get_walkready_joint_command()