diff --git a/simulation/resources/__init__.py b/simulation/resources/__init__.py new file mode 100644 index 0000000..a7a85d1 --- /dev/null +++ b/simulation/resources/__init__.py @@ -0,0 +1,5 @@ +import os + + +def getDataPath(): + return os.path.join(os.path.dirname(__file__)) diff --git a/simulation/resources/spot.urdf b/simulation/resources/spot.urdf new file mode 100644 index 0000000..b4a5492 --- /dev/null +++ b/simulation/resources/spot.urdf @@ -0,0 +1,610 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/simulation/resources/stl/backpart.stl b/simulation/resources/stl/backpart.stl new file mode 100644 index 0000000..375dad5 Binary files /dev/null and b/simulation/resources/stl/backpart.stl differ diff --git a/simulation/resources/stl/foot.stl b/simulation/resources/stl/foot.stl new file mode 100644 index 0000000..ff94767 Binary files /dev/null and b/simulation/resources/stl/foot.stl differ diff --git a/simulation/resources/stl/frontpart.stl b/simulation/resources/stl/frontpart.stl new file mode 100644 index 0000000..536217f Binary files /dev/null and b/simulation/resources/stl/frontpart.stl differ diff --git a/simulation/resources/stl/larm.stl b/simulation/resources/stl/larm.stl new file mode 100644 index 0000000..33e1ced Binary files /dev/null and b/simulation/resources/stl/larm.stl differ diff --git a/simulation/resources/stl/larm_cover.stl b/simulation/resources/stl/larm_cover.stl new file mode 100644 index 0000000..e9e79c0 Binary files /dev/null and b/simulation/resources/stl/larm_cover.stl differ diff --git a/simulation/resources/stl/lfoot.stl b/simulation/resources/stl/lfoot.stl new file mode 100644 index 0000000..754105a Binary files /dev/null and b/simulation/resources/stl/lfoot.stl differ diff --git a/simulation/resources/stl/lshoulder.stl b/simulation/resources/stl/lshoulder.stl new file mode 100644 index 0000000..36e7893 Binary files /dev/null and b/simulation/resources/stl/lshoulder.stl differ diff --git a/simulation/resources/stl/mainbody.stl b/simulation/resources/stl/mainbody.stl new file mode 100644 index 0000000..448fe6e Binary files /dev/null and b/simulation/resources/stl/mainbody.stl differ diff --git a/simulation/resources/stl/rarm.stl b/simulation/resources/stl/rarm.stl new file mode 100644 index 0000000..4fd1eb6 Binary files /dev/null and b/simulation/resources/stl/rarm.stl differ diff --git a/simulation/resources/stl/rarm_cover.stl b/simulation/resources/stl/rarm_cover.stl new file mode 100644 index 0000000..9b5a6e4 Binary files /dev/null and b/simulation/resources/stl/rarm_cover.stl differ diff --git a/simulation/resources/stl/rfoot.stl b/simulation/resources/stl/rfoot.stl new file mode 100644 index 0000000..8c4d3c5 Binary files /dev/null and b/simulation/resources/stl/rfoot.stl differ diff --git a/simulation/resources/stl/rplidar_main.STL b/simulation/resources/stl/rplidar_main.STL new file mode 100644 index 0000000..beced16 Binary files /dev/null and b/simulation/resources/stl/rplidar_main.STL differ diff --git a/simulation/resources/stl/rshoulder.stl b/simulation/resources/stl/rshoulder.stl new file mode 100644 index 0000000..20d7f66 Binary files /dev/null and b/simulation/resources/stl/rshoulder.stl differ diff --git a/simulation/simulation/__init__.py b/simulation/simulation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/simulation/simulation/environment.py b/simulation/simulation/environment.py new file mode 100644 index 0000000..4901845 --- /dev/null +++ b/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 diff --git a/simulation/simulation/robot.py b/simulation/simulation/robot.py new file mode 100644 index 0000000..8dee7a2 --- /dev/null +++ b/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, + ) diff --git a/simulation/test.py b/simulation/test.py new file mode 100644 index 0000000..1c39574 --- /dev/null +++ b/simulation/test.py @@ -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() diff --git a/simulation/train.py b/simulation/train.py new file mode 100644 index 0000000..6136ccf --- /dev/null +++ b/simulation/train.py @@ -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() diff --git a/simulation/training/__init__.py b/simulation/training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/simulation/training/model.py b/simulation/training/model.py new file mode 100644 index 0000000..0660483 --- /dev/null +++ b/simulation/training/model.py @@ -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 diff --git a/simulation/training/trainer.py b/simulation/training/trainer.py new file mode 100644 index 0000000..3a01398 --- /dev/null +++ b/simulation/training/trainer.py @@ -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)