45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
import copy
|
|
from os import sys
|
|
|
|
sys.path.append('../../')
|
|
|
|
from .util.gui import GUI
|
|
|
|
|
|
class Simulator():
|
|
def __init__(self, env, spot) -> None:
|
|
self.env = env
|
|
self.state = env.reset()
|
|
self.gui = GUI(env.spot.quadruped)
|
|
self.action = env.action_space.sample()
|
|
|
|
self.spot = spot
|
|
|
|
def step(self, dt):
|
|
# handle_inputs()
|
|
self.gui.UserInput(self.spot.body_state, self.spot.gait_state)
|
|
self.spot.gait_state.contacts = self.state[-4:]
|
|
self.spot.body_state.worldFeetPositions = copy.deepcopy(self.spot.kinematic.WorldToFoot)
|
|
|
|
self.spot.gait_state.update_gait_state(dt)
|
|
|
|
# Generate next feet positions
|
|
self.spot.gait_controller.generate_trajectory(self.spot.body_state, self.spot.gait_state, dt)
|
|
|
|
# Calculate next joint angles
|
|
joint_angles = self.spot.joint_angles()
|
|
|
|
# Update environment
|
|
self.env.pass_joint_angles(joint_angles.reshape(-1))
|
|
self.state, _, done, _ = self.env.step(self.action)
|
|
if done:
|
|
print("DONE")
|
|
return True
|
|
|
|
def run(self, dt):
|
|
while True:
|
|
done = self.step(dt)
|
|
if done:
|
|
self.env.close()
|
|
|