control updates (#104)

* working mvp

* working setup

sim/envs/humanoids/stompymicro_config.py

@@ -43,9 +43,7 @@ class asset(LeggedRobotCfg.asset):
         termination_height = 0.05
         default_feet_height = 0.01
 
-        terminate_after_contacts_on = [
-            "torso"
-        ]
+        terminate_after_contacts_on = ["torso"]
 
         penalize_contacts_on = []
         self_collisions = 1  # 1 to disable, 0 to enable...bitwise filter
@@ -156,14 +154,14 @@ class rewards:
         target_feet_height = 0.02  # m
         cycle_time = 0.2  # sec
         # if true negative total rewards are clipped at zero (avoids early termination problems)
-        only_positive_rewards = True 
+        only_positive_rewards = True
         # tracking reward = exp(error*sigma)
         tracking_sigma = 5.0
         max_contact_force = 50  # forces above this value are penalized
 
         class scales:
-            
-            # TODO: add an argument 
+
+            # TODO: add an argument
             walking = False
 
             if walking == True:

sim/resources/stompymicro/joints.py

@@ -107,7 +107,6 @@ def default_standing(cls) -> Dict[str, float]:
             cls.legs.right.hip_roll: 0,
             cls.legs.right.knee_pitch: 0,
             cls.legs.right.ankle_pitch: 0,
-
             # TODO: fixing this for debugging
             # Arms
             # cls.left_arm.shoulder_pitch: 0,

sim/resources/stompymicro/robot_fixed.xml

@@ -123,19 +123,7 @@
   </worldbody>
 
   <actuator>
-    <position kp="5" kv="0.3" name="left_hip_pitch" joint="left_hip_pitch" gear="1" forcerange="-23.7 23.7" ctrlrange="-30 30"/>
-    <position kp="5" kv="0.3" name="left_hip_yaw" joint="left_hip_yaw" gear="1" forcerange="-23.7 23.7" ctrlrange="-30 30"/>
-    <position kp="5" kv="0.3" name="left_hip_roll" joint="left_hip_roll" gear="1" forcerange="-23.7 23.7" ctrlrange="-30 30"/>
-    <position kp="5" kv="0.3" name="left_knee_pitch" joint="left_knee_pitch" gear="1" forcerange="-23.7 23.7" ctrlrange="-30 30"/>
-    <position kp="5" kv="0.3" name="left_ankle_pitch" joint="left_ankle_pitch" gear="1" forcerange="-23.7 23.7" ctrlrange="-30 30"/>
-
-    <position kp="5" kv="0.3" name="right_hip_pitch" joint="right_hip_pitch" gear="1" forcerange="-23.7 23.7" ctrlrange="-30 30"/>
-    <position kp="5" kv="0.3" name="right_hip_yaw" joint="right_hip_yaw" gear="1" forcerange="-23.7 23.7" ctrlrange="-30 30"/>
-    <position kp="5" kv="0.3" name="right_hip_roll" joint="right_hip_roll" gear="1" forcerange="-23.7 23.7" ctrlrange="-30 30"/>
-    <position kp="5" kv="0.3" name="right_knee_pitch" joint="right_knee_pitch" gear="1" forcerange="-23.7 23.7" ctrlrange="-30 30"/>
-    <position kp="5" kv="0.3" name="right_ankle_pitch" joint="right_ankle_pitch" gear="1" forcerange="-23.7 23.7" ctrlrange="-30 30"/>
-
-    <!-- <motor name="left_hip_pitch" joint="left_hip_pitch" ctrllimited="true" ctrlrange="-20 20" gear="1" />
+    <motor name="left_hip_pitch" joint="left_hip_pitch" ctrllimited="true" ctrlrange="-20 20" gear="1" />
     <motor name="left_hip_yaw" joint="left_hip_yaw" ctrllimited="true" ctrlrange="-20 20" gear="1" />
     <motor name="left_hip_roll" joint="left_hip_roll" ctrllimited="true" ctrlrange="-20 20" gear="1" />
     <motor name="left_knee_pitch" joint="left_knee_pitch" ctrllimited="true" ctrlrange="-20 20" gear="1" />
@@ -145,11 +133,11 @@
     <motor name="right_hip_yaw" joint="right_hip_yaw" ctrllimited="true" ctrlrange="-20 20" gear="1" />
     <motor name="right_hip_roll" joint="right_hip_roll" ctrllimited="true" ctrlrange="-20 20" gear="1" />
     <motor name="right_knee_pitch" joint="right_knee_pitch" ctrllimited="true" ctrlrange="-20 20" gear="1" />
-    <motor name="right_ankle_pitch" joint="right_ankle_pitch" ctrllimited="true" ctrlrange="-20 20" gear="1" /> -->
+    <motor name="right_ankle_pitch" joint="right_ankle_pitch" ctrllimited="true" ctrlrange="-20 20" gear="1" />
   </actuator>
 
   <sensor>
-    <!-- <actuatorpos name="right_hip_pitch_p" actuator="right_hip_pitch" />
+    <actuatorpos name="right_hip_pitch_p" actuator="right_hip_pitch" />
     <actuatorvel name="right_hip_pitch_v" actuator="right_hip_pitch" />
     <actuatorfrc name="right_hip_pitch_f" actuator="right_hip_pitch" noise="0.001" />
     <actuatorpos name="right_hip_yaw_p" actuator="right_hip_yaw" />
@@ -178,7 +166,7 @@
     <actuatorfrc name="left_knee_pitch_f" actuator="left_knee_pitch" noise="0.001" />
     <actuatorpos name="left_ankle_pitch_p" actuator="left_ankle_pitch" />
     <actuatorvel name="left_ankle_pitch_v" actuator="left_ankle_pitch" />
-    <actuatorfrc name="left_ankle_pitch_f" actuator="left_ankle_pitch" noise="0.001" /> -->
+    <actuatorfrc name="left_ankle_pitch_f" actuator="left_ankle_pitch" noise="0.001" />
     <framequat name="orientation" objtype="site" noise="0.001" objname="imu" />
     <gyro name="angular-velocity" site="imu" noise="0.005" cutoff="34.9" />
   </sensor>

sim/resources/stompypro/joints.py

@@ -67,12 +67,14 @@ class RightLeg(Node):
     knee_pitch = "R_knee"
     ankle_pitch = "R_ankle_y"
 
+
 class LeftArm(Node):
     shoulder_pitch = "L_shoulder_y"
     shoulder_roll = "L_shoulder_z"
     shoulder_yaw = "L_shoulder_x"
     elbow_pitch = "L_elbow_x"
 
+
 class RightArm(Node):
     shoulder_pitch = "R_shoulder_y"
     shoulder_roll = "R_shoulder_z"
@@ -84,10 +86,12 @@ class Legs(Node):
     left = LeftLeg()
     right = RightLeg()
 
+
 class Arms(Node):
     left = LeftArm()
     right = RightArm()
 
+
 class Robot(Node):
     legs = Legs()
     # arms = Arms()
@@ -119,7 +123,7 @@ def default_standing(cls) -> Dict[str, float]:
         return {
             Robot.legs.left.hip_pitch: -0.23,
             Robot.legs.left.hip_yaw: 0.0,
-            Robot.legs.left.hip_roll: 0.0 ,
+            Robot.legs.left.hip_roll: 0.0,
             Robot.legs.left.knee_pitch: 0.441,
             Robot.legs.left.ankle_pitch: -0.258,
             Robot.legs.right.hip_pitch: -0.23,
@@ -143,7 +147,7 @@ def stiffness(cls) -> Dict[str, float]:
             "knee": 300,
             "ankle_y": 300,
         }
-    
+
     # d_gains
     @classmethod
     def damping(cls) -> Dict[str, float]:
@@ -169,11 +173,7 @@ def effort(cls) -> Dict[str, float]:
     # # vel_limits
     @classmethod
     def velocity(cls) -> Dict[str, float]:
-        return {
-            "hip": 10,
-            "knee": 10,
-            "ankle": 10
-        }
+        return {"hip": 10, "knee": 10, "ankle": 10}
 
     @classmethod
     def friction(cls) -> Dict[str, float]:

sim/scripts/create_fixed_torso.py

@@ -52,9 +52,7 @@ def update_urdf(model_path: str, embodiment: str) -> None:
 
 def main() -> None:
     parser = argparse.ArgumentParser(description="Update URDF file to fix robot joints.")
-    parser.add_argument(
-        "--model_path", type=str, help="Path to the model directory", default="sim/resources/stompypro"
-    )
+    parser.add_argument("--model_path", type=str, help="Path to the model directory", default="sim/resources/stompypro")
     parser.add_argument("--embodiment", type=str, help="Embodiment to use", default="stompypro")
     args = parser.parse_args()
 

sim/sim2sim.py

@@ -42,6 +42,7 @@
 import mujoco
 import mujoco_viewer
 import numpy as np
+import pygame
 from scipy.spatial.transform import Rotation as R
 from tqdm import tqdm
 
@@ -50,6 +51,28 @@
 import torch  # isort: skip
 
 
+def handle_keyboard_input():
+    global x_vel_cmd, y_vel_cmd, yaw_vel_cmd
+
+    keys = pygame.key.get_pressed()
+
+    # Update movement commands based on arrow keys
+    if keys[pygame.K_UP]:
+        x_vel_cmd += 0.0005
+    if keys[pygame.K_DOWN]:
+        x_vel_cmd -= 0.0005
+    if keys[pygame.K_LEFT]:
+        y_vel_cmd += 0.0005
+    if keys[pygame.K_RIGHT]:
+        y_vel_cmd -= 0.0005
+
+    # Yaw control
+    if keys[pygame.K_a]:
+        yaw_vel_cmd += 0.001
+    if keys[pygame.K_z]:
+        yaw_vel_cmd -= 0.001
+
+
 class Sim2simCfg:
     def __init__(
         self,
@@ -102,12 +125,6 @@ def __init__(
         self.action_scale = action_scale
 
 
-class cmd:
-    vx = 0.4
-    vy = 0.0
-    dyaw = 0.0
-
-
 def quaternion_to_euler_array(quat):
     # Ensure quaternion is in the correct format [x, y, z, w]
     x, y, z, w = quat
@@ -148,7 +165,7 @@ def pd_control(target_q, q, kp, target_dq, dq, kd, default):
     return kp * (target_q + default - q) - kd * dq
 
 
-def run_mujoco(policy, cfg):
+def run_mujoco(policy, cfg, keyboard_use=False):
     """
     Run the Mujoco simulation using the provided policy and configuration.
 
@@ -192,6 +209,9 @@ def run_mujoco(policy, cfg):
     count_lowlevel = 0
 
     for _ in tqdm(range(int(cfg.sim_duration / cfg.dt)), desc="Simulating..."):
+        if keyboard_use:
+            handle_keyboard_input()
+
         # Obtain an observation
         q, dq, quat, v, omega, gvec = get_obs(data)
         q = q[-cfg.num_actions :]
@@ -209,9 +229,9 @@ def run_mujoco(policy, cfg):
 
             obs[0, 0] = math.sin(2 * math.pi * count_lowlevel * cfg.dt / cfg.cycle_time)
             obs[0, 1] = math.cos(2 * math.pi * count_lowlevel * cfg.dt / cfg.cycle_time)
-            obs[0, 2] = cmd.vx * cfg.lin_vel
-            obs[0, 3] = cmd.vy * cfg.lin_vel
-            obs[0, 4] = cmd.dyaw * cfg.ang_vel
+            obs[0, 2] = x_vel_cmd * cfg.lin_vel
+            obs[0, 3] = y_vel_cmd * cfg.lin_vel
+            obs[0, 4] = yaw_vel_cmd * cfg.ang_vel
             obs[0, 5 : (cfg.num_actions + 5)] = cur_pos_obs
             obs[0, (cfg.num_actions + 5) : (2 * cfg.num_actions + 5)] = cur_vel_obs
             obs[0, (2 * cfg.num_actions + 5) : (3 * cfg.num_actions + 5)] = action
@@ -228,7 +248,6 @@ def run_mujoco(policy, cfg):
                 policy_input[0, i * cfg.num_single_obs : (i + 1) * cfg.num_single_obs] = hist_obs[i][0, :]
 
             action[:] = get_policy_output(policy, policy_input)
-            # action[:] = policy(torch.tensor(policy_input))[0].detach().numpy()
             action = np.clip(action, -cfg.clip_actions, cfg.clip_actions)
             target_q = action * cfg.action_scale
 
@@ -238,7 +257,9 @@ def run_mujoco(policy, cfg):
         tau = pd_control(target_q, q, cfg.kps, target_dq, dq, cfg.kds, default)  # Calc torques
 
         tau = np.clip(tau, -cfg.tau_limit, cfg.tau_limit)  # Clamp torques
-        print(tau)
+        # print(tau)
+        # print(eu_ang)
+        print(x_vel_cmd, y_vel_cmd, yaw_vel_cmd)
 
         data.ctrl = tau
 
@@ -255,8 +276,16 @@ def run_mujoco(policy, cfg):
     parser.add_argument("--embodiment", type=str, required=True, help="embodiment")
     parser.add_argument("--terrain", action="store_true", help="terrain or plane")
     parser.add_argument("--load_actions", action="store_true", help="saved_actions")
+    parser.add_argument("--keyboard_use", action="store_true", help="keyboard_use")
     args = parser.parse_args()
-
+
+    if args.keyboard_use:
+        x_vel_cmd, y_vel_cmd, yaw_vel_cmd = 0.0, 0.0, 0.0
+        pygame.init()
+        pygame.display.set_caption("Simulation Control")
+    else:
+        x_vel_cmd, y_vel_cmd, yaw_vel_cmd = 0.4, 0.0, 0.0
+
     if "pt" in args.load_model:
         policy = torch.jit.load(args.load_model)
     elif "onnx" in args.load_model:
@@ -278,7 +307,7 @@ def get_policy_output(policy, input_data):
             dt=0.001,
             decimation=10,
             cycle_time=0.4,
-            tau_factor=2,
+            tau_factor=3.0,
         )
     elif args.embodiment == "stompymicro":
         cfg = Sim2simCfg(
@@ -290,4 +319,4 @@ def get_policy_output(policy, input_data):
             tau_factor=2,
         )
 
-    run_mujoco(policy, cfg)
+    run_mujoco(policy, cfg, args.keyboard_use)