♻️ Update sim structure

simulation/src/envs/robot.py

@@ -0,0 +1,41 @@
+import pybullet as p
+import numpy as np
+
+
+class QuadrupedRobot:
+    def __init__(self, urdf_path):
+        self.urdf_path = urdf_path
+        self.load()
+
+    def load(self):
+        position = [0, 0, 0.3]
+        orientation = p.getQuaternionFromEuler([0, 0, 0])
+        self.robot_id = p.loadURDF(self.urdf_path, position, orientation)
+
+    def get_observation(self):
+        _, orientation = p.getBasePositionAndOrientation(self.robot_id)
+        orientation = p.getEulerFromQuaternion(orientation)[:2]
+        velocity, angular_velocity = p.getBaseVelocity(self.robot_id)
+        joint_states = p.getJointStates(
+            self.robot_id, range(p.getNumJoints(self.robot_id))
+        )
+        joint_positions = [state[0] for state in joint_states]
+        joint_velocities = [state[1] for state in joint_states]
+        return np.concatenate(
+            [
+                orientation,
+                velocity,
+                angular_velocity,
+                joint_positions,
+                joint_velocities,
+            ]
+        )
+
+    def apply_action(self, action):
+        for i, position in enumerate(action):
+            p.setJointMotorControl2(
+                bodyIndex=self.robot_id,
+                jointIndex=i,
+                controlMode=p.POSITION_CONTROL,
+                targetPosition=position,
+            )