🤖 Adds initial sim structure

This commit is contained in:
Rune Harlyk
2024-08-09 22:27:48 +02:00
committed by Rune Harlyk
parent 2face72aee
commit 33e7fac74c
23 changed files with 798 additions and 0 deletions
+5
View File
@@ -0,0 +1,5 @@
import os
def getDataPath():
return os.path.join(os.path.dirname(__file__))
+610
View File
@@ -0,0 +1,610 @@
<?xml version="1.0" encoding="utf-8"?>
<robot name="rex" xmlns:xacro="http://www.ros.org/wiki/xacro">
<material name="yellow">
<color rgba="0.92 0.83 0.0 1"/>
</material>
<material name="black">
<color rgba="0.1 0.1 0.1 1"/>
</material>
<material name="grey">
<color rgba="0.6 0.6 0.6 1"/>
</material>
<!-- Params -->
<!-- Macros -->
<!-- Robot Body -->
<link name="base_link">
<visual>
<geometry>
<mesh filename="stl/mainbody.stl" scale="0.001 0.001 0.001"/>
</geometry>
<material name="black"/>
<origin rpy="0 0 0" xyz="-0.045 -0.060 -0.015"/>
</visual>
<collision>
<geometry>
<box size="0.14 0.11 0.07"/>
</geometry>
<origin rpy="0 0 0" xyz="0 0 0"/>
</collision>
<inertial>
<mass value="1.20"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<!-- @TODO add lidar sensor-->
<!-- Lidar Backpack -->
<!-- <link name="lidar_link">-->
<!-- <visual>-->
<!-- <geometry>-->
<!-- <mesh filename="stl/rplidar_main.STL" scale="0.001 0.001 0.001"/>-->
<!-- </geometry>-->
<!-- <material name="black"/>-->
<!-- <origin rpy="1.57075 0 0" xyz="-0.06425 0.02885 0"/>-->
<!-- </visual>-->
<!-- <collision>-->
<!-- <geometry>-->
<!-- <box size="0.0985 0.0577 0.07"/>-->
<!-- </geometry>-->
<!-- <origin rpy="0 0 0" xyz="0 0 0"/>-->
<!-- </collision>-->
<!-- <inertial>-->
<!-- <mass value="0.50"/>-->
<!-- <inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>-->
<!-- </inertial>-->
<!-- </link>-->
<!-- <joint name="base_lidar" type="fixed">-->
<!-- <parent link="base_link"/>-->
<!-- <child link="lidar_link"/>-->
<!-- <origin xyz="0 0 0.035"/>-->
<!-- </joint>-->
<!-- create head & tail -->
<link name="chassis_front_link">
<visual>
<geometry>
<mesh filename="stl/frontpart.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.040 -0.060 -0.0140"/>
<material name="yellow"/>
</visual>
<collision>
<geometry>
<box size="0.058 0.11 0.07"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.145 0 0"/>
</collision>
<inertial>
<mass value="0.05"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<joint name="chassis_base_front" type="fixed">
<parent link="base_link"/>
<child link="chassis_front_link"/>
</joint>
<link name="chassis_rear_link">
<visual>
<geometry>
<mesh filename="stl/backpart.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.040 -0.060 -0.0140"/>
<material name="yellow"/>
</visual>
<collision>
<geometry>
<box size="0.04 0.11 0.07"/>
</geometry>
<origin rpy="0 0 0" xyz="0.135 0 0"/>
</collision>
<inertial>
<mass value="0.05"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<joint name="chassis_base_rear" type="fixed">
<parent link="base_link"/>
<child link="chassis_rear_link"/>
</joint>
<!-- create Legs -->
<link name="front_left_shoulder_link">
<visual>
<geometry>
<mesh filename="stl/lshoulder.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.135 -0.02 -0.01"/>
<material name="black"/>
</visual>
<collision>
<geometry>
<box size="0.044 0.038 0.07"/>
</geometry>
<origin rpy="0 0 0" xyz="0 0 0"/>
</collision>
<inertial>
<mass value="0.10"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="front_left_leg_link_cover">
<visual>
<geometry>
<mesh filename="stl/larm_cover.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.125 0.04 -0.02"/>
<material name="yellow"/>
</visual>
<inertial>
<mass value="0.5"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="front_left_leg_link">
<visual>
<geometry>
<mesh filename="stl/larm.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.125 0.04 -0.02"/>
<material name="black"/>
</visual>
<collision>
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 -0.050"/>
<geometry>
<box size="0.028 0.036 0.12"/>
</geometry>
</collision>
<inertial>
<mass value="0.1"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="front_left_foot_link">
<visual>
<geometry>
<mesh filename="stl/lfoot.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.120 0.04 0.1"/>
<material name="black"/>
</visual>
<collision>
<geometry>
<box size="0.026 0.020 0.115"/>
</geometry>
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 -0.050"/>
</collision>
<inertial>
<mass value="0.1"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="front_left_toe_link">
<visual>
<geometry>
<mesh filename="stl/foot.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 -0.40010 0" xyz="0 -0.01 0"/>
<material name="grey"/>
</visual>
<collision>
<geometry>
<mesh filename="stl/foot.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 -0.40010 0" xyz="0 -0.01 0"/>
<contact_coefficients mu="100.0"/>
</collision>
<inertial>
<mass value="0.005"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<joint name="motor_front_left_shoulder" type="revolute">
<parent link="base_link"/>
<child link="front_left_shoulder_link"/>
<axis xyz="1 0 0"/>
<origin rpy="0 0 0" xyz="-0.093 -0.036 0"/>
<limit effort="100.0" lower="-1.0" upper="1.0" velocity="0.7"/>
<dynamics damping="0.0" friction="0.5"/>
</joint>
<joint name="motor_front_left_leg" type="revolute">
<parent link="front_left_shoulder_link"/>
<child link="front_left_leg_link"/>
<axis xyz="0 1 0"/>
<origin rpy="0 0 0" xyz="0 -0.052 0"/>
<limit effort="100.0" lower="-2.17" upper="0.97" velocity="0.5"/>
<dynamics damping="0.0" friction="0.0"/>
</joint>
<joint name="front_left_leg_cover_joint" type="fixed">
<parent link="front_left_leg_link"/>
<child link="front_left_leg_link_cover"/>
<origin xyz="0 0 0"/>
</joint>
<joint name="foot_motor_front_left" type="revolute">
<parent link="front_left_leg_link"/>
<child link="front_left_foot_link"/>
<axis xyz="0 1 0"/>
<origin rpy="0 0 0" xyz="-0.01 0 -0.12"/>
<limit effort="100.0" lower="-0.1" upper="2.59" velocity="0.5"/>
<dynamics damping="0.0" friction="0.5"/>
</joint>
<joint name="front_left_toe" type="fixed">
<parent link="front_left_foot_link"/>
<child link="front_left_toe_link"/>
<origin xyz="0 0 -0.115"/>
</joint>
<link name="front_right_shoulder_link">
<visual>
<geometry>
<mesh filename="stl/rshoulder.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.135 -0.09 -0.01"/>
<material name="black"/>
</visual>
<collision>
<geometry>
<box size="0.044 0.038 0.07"/>
</geometry>
<origin rpy="0 0 0" xyz="0 0 0"/>
</collision>
<inertial>
<mass value="0.10"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="front_right_leg_link_cover">
<visual>
<geometry>
<mesh filename="stl/rarm_cover.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.125 -0.15 -0.02"/>
<material name="yellow"/>
</visual>
<inertial>
<mass value="0.5"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="front_right_leg_link">
<visual>
<geometry>
<mesh filename="stl/rarm.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.125 -0.15 -0.02"/>
<material name="black"/>
</visual>
<collision>
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 -0.050"/>
<geometry>
<box size="0.028 0.036 0.12"/>
</geometry>
</collision>
<inertial>
<mass value="0.1"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="front_right_foot_link">
<visual>
<geometry>
<mesh filename="stl/rfoot.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.120 -0.15 0.1"/>
<material name="black"/>
</visual>
<collision>
<geometry>
<box size="0.026 0.020 0.115"/>
</geometry>
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 -0.050"/>
</collision>
<inertial>
<mass value="0.1"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="front_right_toe_link">
<visual>
<geometry>
<mesh filename="stl/foot.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 -0.40010 0" xyz="0 -0.01 0"/>
<material name="grey"/>
</visual>
<collision>
<geometry>
<mesh filename="stl/foot.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 -0.40010 0" xyz="0 -0.01 0"/>
<contact_coefficients mu="100.0"/>
</collision>
<inertial>
<mass value="0.005"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<joint name="motor_front_right_shoulder" type="revolute">
<parent link="base_link"/>
<child link="front_right_shoulder_link"/>
<axis xyz="1 0 0"/>
<origin rpy="0 0 0" xyz="-0.093 0.036 0"/>
<limit effort="100.0" lower="-1.0" upper="1.0" velocity="0.7"/>
<dynamics damping="0.0" friction="0.5"/>
</joint>
<joint name="motor_front_right_leg" type="revolute">
<parent link="front_right_shoulder_link"/>
<child link="front_right_leg_link"/>
<axis xyz="0 1 0"/>
<origin rpy="0 0 0" xyz="0 0.052 0"/>
<limit effort="100.0" lower="-2.17" upper="0.97" velocity="0.5"/>
<dynamics damping="0.0" friction="0.0"/>
</joint>
<joint name="front_right_leg_cover_joint" type="fixed">
<parent link="front_right_leg_link"/>
<child link="front_right_leg_link_cover"/>
<origin xyz="0 0 0"/>
</joint>
<joint name="foot_motor_front_right" type="revolute">
<parent link="front_right_leg_link"/>
<child link="front_right_foot_link"/>
<axis xyz="0 1 0"/>
<origin rpy="0 0 0" xyz="-0.01 0 -0.12"/>
<limit effort="100.0" lower="-0.1" upper="2.59" velocity="0.5"/>
<dynamics damping="0.0" friction="0.5"/>
</joint>
<joint name="front_right_toe" type="fixed">
<parent link="front_right_foot_link"/>
<child link="front_right_toe_link"/>
<origin xyz="0 0 -0.115"/>
</joint>
<link name="rear_left_shoulder_link">
<visual>
<geometry>
<mesh filename="stl/lshoulder.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.135 -0.02 -0.01"/>
<material name="black"/>
</visual>
<collision>
<geometry>
<box size="0.044 0.038 0.07"/>
</geometry>
<origin rpy="0 0 0" xyz="0 0 0"/>
</collision>
<inertial>
<mass value="0.10"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="rear_left_leg_link_cover">
<visual>
<geometry>
<mesh filename="stl/larm_cover.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.125 0.04 -0.02"/>
<material name="yellow"/>
</visual>
<inertial>
<mass value="0.5"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="rear_left_leg_link">
<visual>
<geometry>
<mesh filename="stl/larm.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.125 0.04 -0.02"/>
<material name="black"/>
</visual>
<collision>
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 -0.050"/>
<geometry>
<box size="0.028 0.036 0.12"/>
</geometry>
</collision>
<inertial>
<mass value="0.10"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="rear_left_foot_link">
<visual>
<geometry>
<mesh filename="stl/lfoot.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.120 0.04 0.1"/>
<material name="black"/>
</visual>
<collision>
<geometry>
<box size="0.026 0.020 0.115"/>
</geometry>
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 -0.050"/>
</collision>
<inertial>
<mass value="0.1"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="rear_left_toe_link">
<visual>
<geometry>
<mesh filename="stl/foot.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 -0.40010 0" xyz="0 -0.01 0"/>
<material name="grey"/>
</visual>
<collision>
<geometry>
<mesh filename="stl/foot.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 -0.40010 0" xyz="0 -0.01 0"/>
<contact_coefficients mu="100.0"/>
</collision>
<inertial>
<mass value="0.005"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<joint name="motor_rear_left_shoulder" type="revolute">
<parent link="base_link"/>
<child link="rear_left_shoulder_link"/>
<axis xyz="1 0 0"/>
<origin rpy="0 0 0" xyz="0.093 -0.036 0"/>
<limit effort="100.0" lower="-1.0" upper="1.0" velocity="0.7"/>
<dynamics damping="0.0" friction="0.5"/>
</joint>
<joint name="motor_rear_left_leg" type="revolute">
<parent link="rear_left_shoulder_link"/>
<child link="rear_left_leg_link"/>
<axis xyz="0 1 0"/>
<origin rpy="0 0 0" xyz="0 -0.052 0"/>
<limit effort="100.0" lower="-2.17" upper="0.97" velocity="0.5"/>
<dynamics damping="0.0" friction="0.0"/>
</joint>
<joint name="rear_left_leg_cover_joint" type="fixed">
<parent link="rear_left_leg_link"/>
<child link="rear_left_leg_link_cover"/>
<origin xyz="0 0 0"/>
</joint>
<joint name="foot_motor_rear_left" type="revolute">
<parent link="rear_left_leg_link"/>
<child link="rear_left_foot_link"/>
<axis xyz="0 1 0"/>
<origin rpy="0 0 0" xyz="-0.01 0 -0.12"/>
<limit effort="100.0" lower="-0.1" upper="2.59" velocity="0.5"/>
<dynamics damping="0.0" friction="0.5"/>
</joint>
<joint name="rear_left_toe" type="fixed">
<parent link="rear_left_foot_link"/>
<child link="rear_left_toe_link"/>
<origin xyz="0 0 -0.115"/>
</joint>
<link name="rear_right_shoulder_link">
<visual>
<geometry>
<mesh filename="stl/rshoulder.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.135 -0.09 -0.01"/>
<material name="black"/>
</visual>
<collision>
<geometry>
<box size="0.044 0.038 0.07"/>
</geometry>
<origin rpy="0 0 0" xyz="0 0 0"/>
</collision>
<inertial>
<mass value="0.10"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="rear_right_leg_link_cover">
<visual>
<geometry>
<mesh filename="stl/rarm_cover.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.125 -0.15 -0.02"/>
<material name="yellow"/>
</visual>
<inertial>
<mass value="0.5"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="rear_right_leg_link">
<visual>
<geometry>
<mesh filename="stl/rarm.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.125 -0.15 -0.02"/>
<material name="black"/>
</visual>
<collision>
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 -0.050"/>
<geometry>
<box size="0.028 0.036 0.12"/>
</geometry>
</collision>
<inertial>
<mass value="0.10"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="rear_right_foot_link">
<visual>
<geometry>
<mesh filename="stl/rfoot.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 0 0" xyz="-0.120 -0.15 0.1"/>
<material name="black"/>
</visual>
<collision>
<geometry>
<box size="0.026 0.020 0.115"/>
</geometry>
<origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 -0.050"/>
</collision>
<inertial>
<mass value="0.1"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<link name="rear_right_toe_link">
<visual>
<geometry>
<mesh filename="stl/foot.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 -0.40010 0" xyz="0 -0.01 0"/>
<material name="grey"/>
</visual>
<collision>
<geometry>
<mesh filename="stl/foot.stl" scale="0.001 0.001 0.001"/>
</geometry>
<origin rpy="0 -0.40010 0" xyz="0 -0.01 0"/>
<contact_coefficients mu="100.0"/>
</collision>
<inertial>
<mass value="0.005"/>
<inertia ixx="100" ixy="0" ixz="0" iyy="100" iyz="0" izz="100"/>
</inertial>
</link>
<joint name="motor_rear_right_shoulder" type="revolute">
<parent link="base_link"/>
<child link="rear_right_shoulder_link"/>
<axis xyz="1 0 0"/>
<origin rpy="0 0 0" xyz="0.093 0.036 0"/>
<limit effort="100.0" lower="-1.0" upper="1.0" velocity="0.7"/>
<dynamics damping="0.0" friction="0.5"/>
</joint>
<joint name="motor_rear_right_leg" type="revolute">
<parent link="rear_right_shoulder_link"/>
<child link="rear_right_leg_link"/>
<axis xyz="0 1 0"/>
<origin rpy="0 0 0" xyz="0 0.052 0"/>
<limit effort="100.0" lower="-2.17" upper="0.97" velocity="0.5"/>
<dynamics damping="0.0" friction="0.0"/>
</joint>
<joint name="rear_right_leg_cover_joint" type="fixed">
<parent link="rear_right_leg_link"/>
<child link="rear_right_leg_link_cover"/>
<origin xyz="0 0 0"/>
</joint>
<joint name="foot_motor_rear_right" type="revolute">
<parent link="rear_right_leg_link"/>
<child link="rear_right_foot_link"/>
<axis xyz="0 1 0"/>
<origin rpy="0 0 0" xyz="-0.01 0 -0.12"/>
<limit effort="100.0" lower="-0.1" upper="2.59" velocity="0.5"/>
<dynamics damping="0.0" friction="0.5"/>
</joint>
<joint name="rear_right_toe" type="fixed">
<parent link="rear_right_foot_link"/>
<child link="rear_right_toe_link"/>
<origin xyz="0 0 -0.115"/>
</joint>
</robot>
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
View File
+36
View File
@@ -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
+42
View File
@@ -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,
)
+24
View File
@@ -0,0 +1,24 @@
from simulation.environment import QuadrupedEnv
from training.model import SimpleNN
import resources as resources
def main():
env = QuadrupedEnv(resources.getDataPath() + "/spot.urdf")
env.reset()
input_size = env.robot.get_observation().shape[0]
output_size = env.robot.get_observation().shape[0]
agent = SimpleNN(input_size, output_size)
done = False
observation = []
while not done:
action = agent.select_action(observation)
observation, reward, done = env.step(action)
if __name__ == "__main__":
main()
+14
View File
@@ -0,0 +1,14 @@
from simulation.environment import QuadrupedEnv
from training.trainer import Trainer
import resources as resources
def main():
env = QuadrupedEnv(resources.getDataPath() + "/spot.urdf")
trainer = Trainer(env)
trainer.train()
if __name__ == "__main__":
main()
View File
+16
View File
@@ -0,0 +1,16 @@
import torch
import torch.nn as nn
class SimpleNN(nn.Module):
def __init__(self, input_size, output_size):
super(SimpleNN, self).__init__()
self.fc1 = nn.Linear(input_size, 128)
self.fc2 = nn.Linear(128, 128)
self.fc3 = nn.Linear(128, output_size)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
x = self.fc3(x)
return x
+51
View File
@@ -0,0 +1,51 @@
import torch
import torch.optim as optim
import pybullet as p
import numpy as np
from tqdm import tqdm, trange
from collections import namedtuple
from training.model import SimpleNN
Experience = namedtuple("Experience", ["observation", "action", "reward", "log_prob"])
class Trainer:
def __init__(self, env):
self.env = env
self.model = SimpleNN(
input_size=env.robot.get_observation().shape[0],
output_size=p.getNumJoints(env.robot.robot_id),
)
self.optimizer = optim.Adam(self.model.parameters(), lr=0.001)
def train(self, episodes=1000):
for episode in trange(episodes):
observation = self.env.reset()
done = False
total_reward = 0
while not done:
action = self.select_action(observation)
observation, reward, done = self.env.step(action)
total_reward += reward
# Train the neural network
# loss = self.compute_loss(observation, action, reward)
# self.optimizer.zero_grad()
# loss.backward()
# self.optimizer.step()
print(f"Episode {episode}: Total Reward: {total_reward}")
def select_action(self, observation):
with torch.no_grad():
observation_tensor = torch.tensor(observation, dtype=torch.float32)
action = self.model(observation_tensor)
return np.array(
[-0.4, -1.5, 6, 0.4, -1.5, 6, -0.4, -1.5, 6, 0.4, -1.5, 6]
) # action.numpy()
def compute_loss(self, observation, action, reward):
# Define your loss function here
return torch.tensor(0.0)