Rune Harlyk
42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
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,
|
|
)
|