🤖 Adds initial sim structure

simulation/simulation/environment.py

@@ -0,0 +1,36 @@
+import pybullet as p
+import pybullet_data
+import numpy as np
+from simulation.robot import QuadrupedRobot
+
+
+class QuadrupedEnv:
+    def __init__(self, urdf_path):
+        p.connect(p.GUI)
+        p.setAdditionalSearchPath(pybullet_data.getDataPath())
+        p.setGravity(0, 0, -9.8)
+        p.setTimeStep(1 / 240)
+
+        self.robot = QuadrupedRobot(urdf_path)
+        self.reset()
+
+    def reset(self):
+        p.resetSimulation()
+        self.robot.load()
+        return self.robot.get_observation()
+
+    def step(self, action):
+        self.robot.apply_action(action)
+        p.stepSimulation()
+        obs = self.robot.get_observation()
+        reward = self.calculate_reward(obs)
+        done = self.is_done(obs)
+        return obs, reward, done
+
+    def calculate_reward(self, observation):
+        # Define your reward function here
+        return 0
+
+    def is_done(self, observation):
+        roll, pitch = observation[0:2]
+        return abs(roll) > 0.5 or abs(pitch) > 0.5

simulation/simulation/robot.py

@@ -0,0 +1,42 @@
+import pybullet as p
+import numpy as np
+
+
+class QuadrupedRobot:
+    def __init__(self, urdf_path):
+        self.urdf_path = urdf_path
+        self.robot_id = None
+
+    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):
+        print(action)
+        for i, position in enumerate(action):
+            p.setJointMotorControl2(
+                bodyIndex=self.robot_id,
+                jointIndex=i,
+                controlMode=p.POSITION_CONTROL,
+                targetPosition=position,
+            )