adjustments to wolfgang_walk_ppo model

Author: sWintermoor

src/bitbots_rl_walk/bitbots_rl_walk/walk.py

@@ -29,7 +29,7 @@
 
 from bitbots_msgs.msg import JointCommand
 
-ONNX_MODEL = os.path.join(get_package_share_directory("bitbots_rl_walk"), "models", "wolfgang_policy.onnx")
+ONNX_MODEL = os.path.join(get_package_share_directory("bitbots_rl_walk"), "models", "wolfgang_walk_ppo.onnx")
 
 PREPARATION_STATE = np.array(
     [
@@ -225,7 +225,8 @@ def _timer_callback(self):
             if self._joint_state is None:
                 self.get_logger().warning("Waiting for joint state data", throttle_duration_sec=1.0)
             if self._cmd_vel is None:
-                self.get_logger().warning("Waiting for cmd_vel data", throttle_duration_sec=1.0)
+                # self.get_logger().warning("Waiting for cmd_vel data", throttle_duration_sec=1.0)
+                self._cmd_vel = Twist(x=1.0, y=0.0, z=0.0)  # Testing purpose
 
             return
 
@@ -275,19 +276,19 @@ def _timer_callback(self):
 
         obs = np.hstack(
             [
-                gyro,
-                gravity,
-                command,
-                joint_angles,
-                joint_velocities,
-                self._previous_action,  # Previous action
-                phase,
+                gyro,  # 3
+                gravity,  # 4
+                command,  # 3
+                joint_angles,  # 18
+                joint_velocities,  # 18
+                self._previous_action,  # 18  # Previous action
+                phase,  # 2
             ]
         ).astype(np.float32)
 
         # Run the ONNX model
-        onnx_input = {"obs": obs.reshape(1, -1)}
-        onnx_pred = self._onnx_session.run(["continuous_actions"], onnx_input)[0][0]
+        onnx_input = {"in_0": obs.reshape(1, -1)}
+        onnx_pred = self._onnx_session.run(["tanh_out_0"], onnx_input)[0][0]
         self._previous_action = onnx_pred
 
         # Publish the joint commands

src/bitbots_rl_walk/setup.py

@@ -0,0 +1,26 @@
+from setuptools import find_packages, setup
+
+package_name = "bitbots_rl_walk"
+
+setup(
+    name=package_name,
+    version="0.0.0",
+    packages=find_packages(exclude=["test"]),
+    data_files=[
+        ("share/ament_index/resource_index/packages", ["resource/" + package_name]),
+        ("share/" + package_name, ["package.xml"]),
+        ("share/" + package_name + "/models", ["models/wolfgang_walk_ppo.onnx"]),
+    ],
+    install_requires=["setuptools"],
+    zip_safe=True,
+    maintainer="florian",
+    maintainer_email="git@flova.de",
+    description="TODO: Package description",
+    license="TODO: License declaration",
+    tests_require=["pytest"],
+    entry_points={
+        "console_scripts": [
+            "walk = bitbots_rl_walk.walk:main",
+        ],
+    },
+)