diff --git a/sim/macro.py b/sim/macro.py new file mode 100644 index 00000000..3e2cac05 --- /dev/null +++ b/sim/macro.py @@ -0,0 +1,236 @@ +""" +python sim/macro.py --embodiment stompymicro +""" +import argparse +import math +import os +from collections import deque +from copy import deepcopy + +import mujoco +import mujoco_viewer +import numpy as np +from scipy.spatial.transform import Rotation as R +from tqdm import tqdm + +from sim.scripts.create_mjcf import load_embodiment + + +class Sim2simCfg: + def __init__( + self, + embodiment, + frame_stack=15, + c_frame_stack=3, + sim_duration=60.0, + dt=0.001, + decimation=10, + cycle_time=0.4, + tau_factor=3, + lin_vel=2.0, + ang_vel=1.0, + dof_pos=1.0, + dof_vel=0.05, + clip_observations=18.0, + clip_actions=18.0, + action_scale=0.25, + ): + self.robot = load_embodiment(embodiment) + + self.num_actions = len(self.robot.all_joints()) + self.frame_stack = frame_stack + self.c_frame_stack = c_frame_stack + self.num_single_obs = 11 + self.num_actions * self.c_frame_stack + self.num_observations = int(self.frame_stack * self.num_single_obs) + + self.sim_duration = sim_duration + self.dt = dt + self.decimation = decimation + + self.cycle_time = cycle_time + self.kps = np.array(list(self.robot.stiffness().values()) + list(self.robot.stiffness().values())) + self.kds = np.array(list(self.robot.damping().values()) + list(self.robot.damping().values())) + + self.lin_vel = lin_vel + self.ang_vel = ang_vel + self.dof_pos = dof_pos + self.dof_vel = dof_vel + + self.clip_observations = clip_observations + self.clip_actions = clip_actions + + self.action_scale = action_scale + + +def pd_control(target_q, q, kp, target_dq, dq, kd, default): + """Calculates torques from position commands""" + return kp * (target_q + default - q) - kd * dq + + +def get_scripted_joint_targets(current_time, cfg, joints): + """ + Returns the target joint positions for the robot at the given time. + + Args: + current_time (float): The current simulation time. + cfg (Sim2simCfg): The simulation configuration. + + Returns: + np.ndarray: The target joint positions. + """ + global save_q + + # Initialize target positions + target_q = np.zeros(cfg.num_actions, dtype=np.double) + + # Define time intervals for different actions + # You can adjust these intervals as needed + t0 = .033 # Time to let it fall + t1 = 1.0 # Time to move arms to position + t2 = 2.0 # Time to bend knees + t3 = 3.0 # Time to stand up + t4 = 4.0 # Time to bend knees + # lie down + # Lie down to arms up (0 to t1) + # breakpoint() + # let it fall + # + # + if current_time <= t0: + pass + # else: + elif current_time <= t1: + progress_1 = (current_time - t0) / (t1 - t0) + # Interpolate arm joints from initial to target position + # Assuming arm joints are indices 0 and 1 (adjust based on your robot) + arm_joint_indices = [joints["right_shoulder_pitch"], joints["left_shoulder_pitch"]] # Replace with actual indices + arm_target_positions = [-1.55, 1.55] # Raise arms up + + for idx, joint in enumerate(arm_joint_indices): + target_q[joint] = np.interp(progress_1, [0, 1], [0, arm_target_positions[idx]]) + + save_q = target_q + elif current_time <= t2: + # else: + target_q = save_q + progress_2 = (current_time - t1) / (t2 - t1) + # Interpolate arm joints from initial to target position + # Assuming arm joints are indices 0 and 1 (adjust based on your robot) + arm_joint_indices = [joints["right_shoulder_yaw"], joints["left_shoulder_yaw"]] # Replace with actual indices + arm_target_positions = [0.436, -0.436, ] # Raise arms up + + for idx, joint in enumerate(arm_joint_indices): + target_q[joint] = np.interp(progress_2, [0, 1], [0, arm_target_positions[idx]]) + + save_q = target_q + # Arms up to bend knees (t1 to t2) + elif current_time <= t3: + target_q = save_q + progress = (current_time - t1) / (t2 - t1) + + knee_joint_indices = [joints["right_knee_pitch"], joints["left_knee_pitch"]] # Replace with actual indices + knee_bend_positions = np.deg2rad([-90, 90]) # Bend knees + for idx, joint in enumerate(knee_joint_indices): + target_q[joint] = np.interp(progress, [0, 1], [0, knee_bend_positions[idx]]) + save_q = target_q + # Bend knees to stand up (t2 to t3) + + else: + target_q = save_q + progress = (current_time - t3) / (t4 - t3) + # Knees extend to standing position + pitch_joint_indices = [joints["right_hip_pitch"], joints["left_hip_pitch"]] # Replace with actual indices + pitch_positions = np.deg2rad([-90, 90]) # Bend knees + for idx, joint in enumerate(pitch_joint_indices): + target_q[joint] = np.interp(progress, [0, 1], [0, pitch_positions[idx]]) + save_q = target_q + + return target_q + +def run_mujoco_scripted(cfg): + """ + Run the Mujoco simulation using scripted joint positions. + + Args: + cfg: The configuration object containing simulation settings. + + Returns: + None + """ + model_dir = os.environ.get("MODEL_DIR") + mujoco_model_path = f"{model_dir}/{args.embodiment}/robot_fixed.xml" + + model = mujoco.MjModel.from_xml_path(mujoco_model_path) + model.opt.timestep = cfg.dt + data = mujoco.MjData(model) + + # Initialize default positions + try: + data.qpos = model.keyframe("default").qpos + default = deepcopy(model.keyframe("default").qpos)[-cfg.num_actions:] + print("Default position:", default) + except: + print("No default position found, using zero initialization") + default = np.zeros(cfg.num_actions) + + mujoco.mj_step(model, data) + data.qvel = np.zeros_like(data.qvel) + data.qacc = np.zeros_like(data.qacc) + viewer = mujoco_viewer.MujocoViewer(model, data) + + mujoco.mj_step(model, data) + joints = {} + for ii in range(1, len(data.ctrl)): + joints[data.joint(ii).name] = data.joint(ii).id - 1 + + # Initialize target joint positions + target_q = np.zeros((cfg.num_actions), dtype=np.double) + + # Simulation loop + sim_steps = int(cfg.sim_duration / cfg.dt) + for step in tqdm(range(sim_steps), desc="Simulating..."): + current_time = step * cfg.dt + q = data.qpos[-cfg.num_actions:].astype(np.double) + dq = data.qvel[-cfg.num_actions:].astype(np.double) + + # Update target_q based on scripted behavior + target_q = get_scripted_joint_targets(current_time, cfg, joints) + target_dq = np.zeros((cfg.num_actions), dtype=np.double) + tau = pd_control(target_q, q, cfg.kps, target_dq, dq, cfg.kds, default) + + print("tau:", tau) + print("q:", q) + print("target_q:", target_q) + data.ctrl = tau + + mujoco.mj_step(model, data) + viewer.render() + + viewer.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Deployment script.") + parser.add_argument("--embodiment", type=str, required=True, help="embodiment") + args = parser.parse_args() + + if args.embodiment == "stompypro": + cfg = Sim2simCfg( + args.embodiment, + sim_duration=60.0, + dt=0.001, + decimation=10, + cycle_time=0.4, + tau_factor=3.0, + ) + elif args.embodiment == "stompymicro": + cfg = Sim2simCfg( + args.embodiment, + sim_duration=60.0, + dt=0.001, + decimation=10, + cycle_time=0.4, + tau_factor=10., + ) + + run_mujoco_scripted(cfg) diff --git a/sim/resources/stompymicro/joints.py b/sim/resources/stompymicro/joints.py index 78156914..4ac62cf2 100644 --- a/sim/resources/stompymicro/joints.py +++ b/sim/resources/stompymicro/joints.py @@ -55,13 +55,13 @@ def __str__(self) -> str: class LeftArm(Node): shoulder_yaw = "left_shoulder_yaw" shoulder_pitch = "left_shoulder_pitch" - elbow_pitch = "left_elbow_yaw" # FIXME: yaw vs pitch + elbow_yaw = "left_elbow_yaw" # FIXME: yaw vs pitch class RightArm(Node): shoulder_yaw = "right_shoulder_yaw" shoulder_pitch = "right_shoulder_pitch" - elbow_pitch = "right_elbow_yaw" # FIXME: yaw vs pitch + elbow_yaw = "right_elbow_yaw" # FIXME: yaw vs pitch class LeftLeg(Node): @@ -89,63 +89,63 @@ class Robot(Node): height = 0.293 rotation = [0, 0, 0.707, 0.707] - # left_arm = LeftArm() - # right_arm = RightArm() + left_arm = LeftArm() + right_arm = RightArm() legs = Legs() @classmethod def default_standing(cls) -> Dict[str, float]: return { # Legs - cls.legs.left.hip_pitch: 0.23, + cls.legs.left.hip_pitch: 0.0, cls.legs.left.hip_yaw: 0, cls.legs.left.hip_roll: 0, - cls.legs.left.knee_pitch: 0.741, - cls.legs.left.ankle_pitch: -0.5, - cls.legs.right.hip_pitch: -0.23, + cls.legs.left.knee_pitch: 0.0, + cls.legs.left.ankle_pitch: -0.0, + cls.legs.right.hip_pitch: -0.0, cls.legs.right.hip_yaw: 0, cls.legs.right.hip_roll: 0, - cls.legs.right.knee_pitch: -0.741, - cls.legs.right.ankle_pitch: 0.5, + cls.legs.right.knee_pitch: 0.0, + cls.legs.right.ankle_pitch: 0.0, # TODO: fixing this for debugging # Arms - # cls.left_arm.shoulder_pitch: 0, - # cls.left_arm.shoulder_yaw: 0, - # cls.left_arm.elbow_pitch: 0, - # cls.right_arm.shoulder_pitch: 0, - # cls.right_arm.shoulder_yaw: 0, - # cls.right_arm.elbow_pitch: 0, + cls.left_arm.shoulder_pitch: 0, + cls.left_arm.shoulder_yaw: 0, + cls.left_arm.elbow_yaw: 0, + cls.right_arm.shoulder_pitch: 0, + cls.right_arm.shoulder_yaw: 0, + cls.right_arm.elbow_yaw: 0, } @classmethod def default_limits(cls) -> Dict[str, Dict[str, float]]: return { - # # left arm - # Robot.left_arm.shoulder_pitch: { - # "lower": -1.5707963, - # "upper": 1.5707963, - # }, - # Robot.left_arm.shoulder_yaw: { - # "lower": -1.5707963, - # "upper": 1.5707963, - # }, - # Robot.left_arm.elbow_pitch: { - # "lower": -1.2217305, - # "upper": 1.2217305, - # }, - # # right arm - # Robot.right_arm.shoulder_pitch: { - # "lower": -1.5707963, - # "upper": 1.5707963, - # }, - # Robot.right_arm.shoulder_yaw: { - # "lower": -1.5707963, - # "upper": 1.5707963, - # }, - # Robot.right_arm.elbow_pitch: { - # "lower": -1.2217305, - # "upper": 1.2217305, - # }, + # left arm + Robot.left_arm.shoulder_pitch: { + "lower": -1.5707963, + "upper": 1.5707963, + }, + Robot.left_arm.shoulder_yaw: { + "lower": -1.5707963, + "upper": 1.5707963, + }, + Robot.left_arm.elbow_yaw: { + "lower": -1.5707963, + "upper": 1.5707963, + }, + # right arm + Robot.right_arm.shoulder_pitch: { + "lower": -1.5707963, + "upper": 1.5707963, + }, + Robot.right_arm.shoulder_yaw: { + "lower": -1.5707963, + "upper": 1.5707963, + }, + Robot.right_arm.elbow_yaw: { + "lower": -1.5707963, + "upper": 1.5707963, + }, # left leg Robot.legs.left.hip_pitch: { "lower": -1.5707963, @@ -194,16 +194,14 @@ def default_limits(cls) -> Dict[str, Dict[str, float]]: @classmethod def stiffness(cls) -> Dict[str, float]: return { - "hip_pitch": 17.681462808698132, - "hip_yaw": 17.681462808698132, - "hip_roll": 17.681462808698132, - "knee_pitch": 17.681462808698132, - "ankle_pitch": 17.681462808698132, - # "shoulder_pitch": 5, - # "shoulder_yaw": 5, - # "shoulder_roll": 5, - # "elbow_pitch": 5, - # "elbow_yaw": 5, + "hip_pitch": 5.5, + "hip_yaw": 5.5, + "hip_roll": 5.5, + "knee_pitch": 5.5, + "ankle_pitch": 5.5, + "shoulder_pitch": 5.5, + "shoulder_yaw": 5.5, + "elbow_yaw": 5.5, } # d_gains @@ -215,11 +213,9 @@ def damping(cls) -> Dict[str, float]: "hip_roll": 0.5354656169048285, "knee_pitch": 0.5354656169048285, "ankle_pitch": 0.5354656169048285, - # "shoulder_pitch": 0.3, - # "shoulder_yaw": 0.3, - # "shoulder_roll": 0.3, - # "elbow_pitch": 0.3, - # "elbow_yaw": 0.3, + "shoulder_pitch": 0.5354656169048285, + "shoulder_yaw": 0.5354656169048285, + "elbow_yaw": 0.5354656169048285, } # pos_limits @@ -231,11 +227,9 @@ def effort(cls) -> Dict[str, float]: "hip_roll": 1, "knee_pitch": 1, "ankle_pitch": 1, - # "shoulder_pitch": 1, - # "shoulder_yaw": 1, - # "shoulder_roll": 1, - # "elbow_pitch": 1, - # "elbow_yaw": 1, + "shoulder_pitch": 1, + "shoulder_yaw": 1, + "elbow_yaw": 1, } # vel_limits @@ -247,11 +241,9 @@ def velocity(cls) -> Dict[str, float]: "hip_roll": 10, "knee_pitch": 10, "ankle_pitch": 10, - # "shoulder_pitch": 10, - # "shoulder_yaw": 10, - # "shoulder_roll": 10, - # "elbow_pitch": 10, - # "elbow_yaw": 10, + "shoulder_pitch": 10, + "shoulder_yaw": 10, + "elbow_yaw": 10, } @classmethod @@ -263,9 +255,7 @@ def friction(cls) -> Dict[str, float]: "hip_roll": 0.05, "knee_pitch": 0.05, "ankle_pitch": 0.05, - # "ankle_pitch": 0.05, - # "elbow_yaw": 0.05, - # "elbow_pitch": 0.05, + "elbow_yaw": 0.05, } diff --git a/sim/resources/stompymicro/robot.xml b/sim/resources/stompymicro/robot.xml new file mode 100644 index 00000000..4f9e8a25 --- /dev/null +++ b/sim/resources/stompymicro/robot.xml @@ -0,0 +1,107 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/sim/resources/stompymicro/robot_calibration.xml b/sim/resources/stompymicro/robot_calibration.xml deleted file mode 100644 index 199afdbf..00000000 --- a/sim/resources/stompymicro/robot_calibration.xml +++ /dev/null @@ -1,181 +0,0 @@ - - - - - - - - - - - - diff --git a/sim/resources/stompymicro/robot_fixed.xml b/sim/resources/stompymicro/robot_fixed.xml index f078e8fd..a025b5aa 100644 --- a/sim/resources/stompymicro/robot_fixed.xml +++ b/sim/resources/stompymicro/robot_fixed.xml @@ -1,119 +1,139 @@ - + \ No newline at end of file diff --git a/sim/scripts/create_fixed_torso.py b/sim/scripts/create_fixed_torso.py index 7002eb9b..3dde8870 100755 --- a/sim/scripts/create_fixed_torso.py +++ b/sim/scripts/create_fixed_torso.py @@ -6,12 +6,13 @@ from pathlib import Path from sim.scripts.create_mjcf import load_embodiment - +from sim.scripts.create_mjcf import create_mjcf def update_urdf(model_path: str, embodiment: str) -> None: tree = ET.parse(Path(model_path) / "robot.urdf") root = tree.getroot() robot = load_embodiment(embodiment) + print(robot.default_standing()) revolute_joints = set(robot.default_standing().keys()) @@ -57,7 +58,7 @@ def main() -> None: args = parser.parse_args() update_urdf(args.model_path, args.embodiment) - # create_mjcf(Path(args.model_path) / "robot") + create_mjcf(Path(args.model_path) / "robot") if __name__ == "__main__":