From 8ef567fad8f4927dd0af5f6f4b75742d614df7c0 Mon Sep 17 00:00:00 2001 From: budzianowski Date: Fri, 12 Jul 2024 01:53:05 -0700 Subject: [PATCH 1/2] add draft of the real policy --- sim/deploy/run.py | 42 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/sim/deploy/run.py b/sim/deploy/run.py index 9f5e02fe..cc3eba39 100755 --- a/sim/deploy/run.py +++ b/sim/deploy/run.py @@ -2,6 +2,8 @@ Run example: mjpython sim/deploy/run.py --load_model sim/deploy/tests/walking_policy.pt --world MUJOCO + + python sim/deploy/run.py --load_model MODEL_WEIGHTS --world REAL """ import argparse @@ -23,6 +25,11 @@ from sim.env import stompy_mjcf_path from sim.stompy.joints import StompyFixed +import time +from firmware.cpp.imu.imu import IMU +from firmware.scripts.robot_controller import Robot # TODO:(ved) move this to a more appropriate location +import torch + class Worlds(Enum): MUJOCO = "SIM" @@ -50,6 +57,41 @@ def get_observation(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarra pass +class Real(World): + def __init__(self, cfg: RobotConfig): + self.robot = Robot("legs") + self.robot.zero_out() # TODO: (Ved - zero out the legs) + self.model = torch.load(cfg.robot_model_path) # TODO: (Allen/Isaac - load the model) + self.imu = IMU(1) # TODO: (Weasley -load the imu) + self.state = None + + def step(self, observation: np.ndarray) -> None: + """Performs a simulation in the real world.""" + tau = self.model(observation) # TODO: (Allen/Isaac - run the model) + self.robot.set_position(tau) # TODO: (Ved - set the position of the robot) + + def get_observation(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Extracts an observation from the world state. + + Returns: + A tuple containing the following: + - dof_pos: The joint positions. + - dof_vel: The joint velocities. + - orientation: The orientation of the robot. + - ang_vel: The angular velocity of the robot. + """ + ang_vel, orientation = self.imu.step() + dof_pos = self.robot.get_position() + dof_vel = self.robot.get_velocity() + return (dof_pos, dof_vel, orientation, ang_vel) + + def simulate(self, policy=None) -> None: + for step in tqdm(range(int(cfg.duration / self.cfg.dt)), desc="Simulating..."): + obs = self.get_observation() + action = self.step(obs) + time.sleep(self.cfg.dt) + + class MujocoWorld(World): """Simulated world using MuJoCo. From 5e2e7125967244b766bc5701bdf01596348d9d93 Mon Sep 17 00:00:00 2001 From: WT-MM Date: Mon, 15 Jul 2024 11:13:41 -0700 Subject: [PATCH 2/2] integrate imu --- sim/deploy/run.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/sim/deploy/run.py b/sim/deploy/run.py index cc3eba39..4acbc575 100755 --- a/sim/deploy/run.py +++ b/sim/deploy/run.py @@ -26,7 +26,7 @@ from sim.stompy.joints import StompyFixed import time -from firmware.cpp.imu.imu import IMU +from firmware.imu.imu import IMUInterface from firmware.scripts.robot_controller import Robot # TODO:(ved) move this to a more appropriate location import torch @@ -62,13 +62,15 @@ def __init__(self, cfg: RobotConfig): self.robot = Robot("legs") self.robot.zero_out() # TODO: (Ved - zero out the legs) self.model = torch.load(cfg.robot_model_path) # TODO: (Allen/Isaac - load the model) - self.imu = IMU(1) # TODO: (Weasley -load the imu) + self.imu = IMUInterface(1) # Bus = 1 self.state = None def step(self, observation: np.ndarray) -> None: """Performs a simulation in the real world.""" tau = self.model(observation) # TODO: (Allen/Isaac - run the model) self.robot.set_position(tau) # TODO: (Ved - set the position of the robot) + self.imu.step(0.01) # Replace with actual dt + def get_observation(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Extracts an observation from the world state. @@ -80,7 +82,7 @@ def get_observation(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarra - orientation: The orientation of the robot. - ang_vel: The angular velocity of the robot. """ - ang_vel, orientation = self.imu.step() + ang_vel, orientation = self.imu.state dof_pos = self.robot.get_position() dof_vel = self.robot.get_velocity() return (dof_pos, dof_vel, orientation, ang_vel)