some stuff

sim/humanoid_gym/envs/humanoid_env.py

@@ -208,11 +208,18 @@ def reset_idx(self, env_ids: Tensor) -> None:
         for i in range(self.critic_history.maxlen):
             self.critic_history[i][env_ids] *= 0
 
+    # def _compute_torques(self, actions: Tensor) -> Tensor:
+    #     # Need to override this so we can just use the motor torques directly.
+    #     # return torch.clip(actions * self.cfg.control.action_scale, -self.torque_limits, self.torque_limits)
+    #     torques = torch.tanh(actions * self.cfg.control.action_scale) * self.torque_limits
+    #     return torques
+
     def _compute_torques(self, actions: Tensor) -> Tensor:
-        # Need to override this so we can just use the motor torques directly.
-        # return torch.clip(actions * self.cfg.control.action_scale, -self.torque_limits, self.torque_limits)
-        torques = torch.tanh(actions * self.cfg.control.action_scale) * self.torque_limits
-        return torques
+        actions_scaled = actions * self.cfg.control.action_scale
+        p_gains = 80.0
+        d_gains = 5.0
+        torques = p_gains * (actions_scaled + self.default_dof_pos - self.dof_pos) - d_gains * self.dof_vel
+        return torch.clip(torques, -self.torque_limits, self.torque_limits)
 
     def check_termination(self) -> None:
         self.reset_buf = self.episode_length_buf > self.max_episode_length

sim/scripts/download_urdf.sh

@@ -17,4 +17,5 @@ kol urdf ${url} \
         dof_x10=12 \
         knee_revolute=13.9 \
         ankle_revolute=6 \
-    --output-dir ${output_dir}
+    --output-dir ${output_dir} \
+    --disable-mimics

sim/scripts/drop_urdf.py

@@ -27,15 +27,18 @@
 Viewer = NewType("Viewer", Any)
 Args = NewType("Args", Any)
 
-DRIVE_MODE = gymapi.DOF_MODE_EFFORT
-# DRIVE_MODE = gymapi.DOF_MODE_POS
+# Importing torch down here to avoid gymtorch issues.
+import torch  # noqa: E402
+
+# DRIVE_MODE = gymapi.DOF_MODE_EFFORT
+DRIVE_MODE = gymapi.DOF_MODE_POS
 
 # Stiffness and damping are Kp and Kd parameters for the PD controller
 # that drives the joints to the desired position.
-# STIFFNESS = 80.0
-# DAMPING = 5.0
-STIFFNESS = 0.0
-DAMPING = 0.0
+STIFFNESS = 80.0
+DAMPING = 5.0
+# STIFFNESS = 0.0
+# DAMPING = 0.0
 
 # Armature is a parameter that can be used to model the inertia of the joint.
 # We set it to zero because the URDF already models the inertia of the joints.
@@ -114,12 +117,15 @@ def load_gym() -> GymParams:
     asset_options = gymapi.AssetOptions()
     asset_options.default_dof_drive_mode = DRIVE_MODE
     asset_options.collapse_fixed_joints = True
+    asset_options.disable_gravity = True
+    asset_options.fix_base_link = True
     asset_path = stompy_urdf_path()
     robot_asset = gym.load_urdf(sim, str(asset_path.parent), str(asset_path.name), asset_options)
 
     # Adds the robot to the environment.
     initial_pose = gymapi.Transform()
     initial_pose.p = gymapi.Vec3(0.0, 5.0, 0.0)
+    initial_pose.r = gymapi.Quat(0.0, 0.0, 0.0, 1.0)
     robot = gym.create_actor(env, robot_asset, initial_pose, "robot")
 
     # Configure DOF properties.
@@ -135,6 +141,28 @@ def load_gym() -> GymParams:
     cam_target = gymapi.Vec3(0, 2, 1.5)
     gym.viewer_camera_look_at(viewer, None, cam_pos, cam_target)
 
+    # Gets tensors for the DOF states.
+    dof_state_tensor = gym.acquire_dof_state_tensor(sim)
+    gym.refresh_dof_state_tensor(sim)
+    dof_state = gymtorch.wrap_tensor(dof_state_tensor)
+    num_dof = len(Stompy.all_joints())
+    dof_pos = dof_state.view(1, num_dof, 2)[..., 0]
+    # dof_vel = dof_state.view(1, num_dof, 2)[..., 1]
+
+    # Resets the DOF positions to the starting positions.
+    # dof_vel[:] = 0.0
+    starting_positions = Stompy.default_positions()
+    dof_ids: Dict[str, int] = gym.get_actor_dof_dict(env, robot)
+    for joint_name, joint_position in starting_positions.items():
+        dof_pos[0, dof_ids[joint_name]] = joint_position
+    env_ids_int32 = torch.zeros(1, dtype=torch.int32)
+    gym.set_dof_state_tensor_indexed(
+        sim,
+        gymtorch.unwrap_tensor(dof_state),
+        gymtorch.unwrap_tensor(env_ids_int32),
+        1,
+    )
+
     return GymParams(
         gym=gym,
         env=env,
@@ -167,36 +195,41 @@ def run_gym(gym: GymParams, mode: Literal["one_at_a_time", "all_at_once"] = "all
         # print(gym.gym.get_env_rigid_contact_forces(gym.env))
 
         # Prints the contact forces.
-        net_contact_forces = gym.gym.acquire_net_contact_force_tensor(gym.sim)
-        net_contact_forces_tensor = gymtorch.wrap_tensor(net_contact_forces).norm(2, dim=-1)  # (38, 3)
-        if net_contact_forces_tensor.size(0) != len(body_ids):
-            raise ValueError("Mismatch between body IDs and contact forces.")
-        for body_id, contact_force in zip(body_ids, net_contact_forces_tensor):
-            contact_force = contact_force.item()
-            logger.info("Body %s: %.3g", body_id, contact_force)
+        # net_contact_forces = gym.gym.acquire_net_contact_force_tensor(gym.sim)
+        # net_contact_forces_tensor = gymtorch.wrap_tensor(net_contact_forces).norm(2, dim=-1)  # (38, 3)
+        # if net_contact_forces_tensor.size(0) != len(body_ids):
+        #     raise ValueError("Mismatch between body IDs and contact forces.")
+        # for body_id, contact_force in zip(body_ids, net_contact_forces_tensor):
+        #     contact_force = contact_force.item()
+        #     logger.info("Body %s: %.3g", body_id, contact_force)
+
+        # Prints the joint angles.
+        # joint_positions = gym.gym.get_actor_dof_states(gym.env, gym.robot, gymapi.STATE_ALL)
+        # for joint_name, (joint_position, joint_velocity) in zip(joints, joint_positions):
+        #     logger.info("Joint %s: %.3g %.3g", joint_name, joint_position, joint_velocity)
 
         # Every second, set the target effort for each joint to the reverse.
-        curr_time = time.time()
-
-        if mode == "one_at_a_time":
-            if curr_time - last_time > 0.25:
-                last_time = curr_time
-                gym.gym.apply_dof_effort(gym.env, joint_id, 0.0)
-                joint_id += 1
-                if joint_id >= len(joints):
-                    effort = -effort
-                    joint_id = 0
-                gym.gym.apply_dof_effort(gym.env, joint_id, effort)
-
-        elif mode == "all_at_once":
-            if curr_time - last_time > 1.0:
-                last_time = curr_time
-                effort = -effort
-                for joint_name in joints:
-                    gym.gym.apply_dof_effort(gym.env, dof_ids[joint_name], effort)
-
-        else:
-            raise ValueError(f"Invalid mode: {mode}")
+        # curr_time = time.time()
+
+        # if mode == "one_at_a_time":
+        #     if curr_time - last_time > 0.25:
+        #         last_time = curr_time
+        #         gym.gym.apply_dof_effort(gym.env, joint_id, 0.0)
+        #         joint_id += 1
+        #         if joint_id >= len(joints):
+        #             effort = -effort
+        #             joint_id = 0
+        #         gym.gym.apply_dof_effort(gym.env, joint_id, effort)
+
+        # elif mode == "all_at_once":
+        #     if curr_time - last_time > 1.0:
+        #         last_time = curr_time
+        #         effort = -effort
+        #         for joint_name in joints:
+        #             gym.gym.apply_dof_effort(gym.env, dof_ids[joint_name], effort)
+
+        # else:
+        #     raise ValueError(f"Invalid mode: {mode}")
 
     gym.gym.destroy_viewer(gym.viewer)
     gym.gym.destroy_sim(gym.sim)

sim/stompy/joints.py

@@ -7,7 +7,9 @@
 
 import textwrap
 from abc import ABC
-from typing import List, Union
+from typing import Dict, List, Union
+
+import numpy as np
 
 
 class Node(ABC):
@@ -42,79 +44,90 @@ def __str__(self) -> str:
 
 
 class Head(Node):
-    left_right: str = "joint_head_1_x4_1_dof_x4"
-    up_down: str = "joint_head_1_x4_2_dof_x4"
+    left_right = "joint_head_1_x4_1_dof_x4"
+    up_down = "joint_head_1_x4_2_dof_x4"
 
 
 class Torso(Node):
-    pitch: str = "joint_torso_1_x8_1_dof_x8"
+    pitch = "joint_torso_1_x8_1_dof_x8"
 
 
 class LeftHand(Node):
-    hand_roll: str = "joint_left_arm_2_hand_1_x4_1_dof_x4"
-    hand_grip: str = "joint_left_arm_2_hand_1_x4_2_dof_x4"
-    slider_a: str = "joint_left_arm_2_hand_1_slider_1"
-    slider_b: str = "joint_left_arm_2_hand_1_slider_2"
+    hand_roll = "joint_left_arm_2_hand_1_x4_1_dof_x4"
+    hand_grip = "joint_left_arm_2_hand_1_x4_2_dof_x4"
+    slider_a = "joint_left_arm_2_hand_1_slider_1"
+    slider_b = "joint_left_arm_2_hand_1_slider_2"
 
 
 class LeftArm(Node):
-    shoulder_yaw: str = "joint_left_arm_2_x8_1_dof_x8"
-    shoulder_pitch: str = "joint_left_arm_2_x8_2_dof_x8"
-    shoulder_roll: str = "joint_left_arm_2_x6_1_dof_x6"
-    elbow_yaw: str = "joint_left_arm_2_x6_2_dof_x6"
-    elbow_roll: str = "joint_left_arm_2_x4_1_dof_x4"
-    hand: Node = LeftHand()
+    shoulder_yaw = "joint_left_arm_2_x8_1_dof_x8"
+    shoulder_pitch = "joint_left_arm_2_x8_2_dof_x8"
+    shoulder_roll = "joint_left_arm_2_x6_1_dof_x6"
+    elbow_yaw = "joint_left_arm_2_x6_2_dof_x6"
+    elbow_roll = "joint_left_arm_2_x4_1_dof_x4"
+    hand = LeftHand()
 
 
 class RightHand(Node):
-    hand_roll: str = "joint_right_arm_1_hand_1_x4_1_dof_x4"
-    hand_grip: str = "joint_right_arm_1_hand_1_x4_2_dof_x4"
-    slider_a: str = "joint_right_arm_1_hand_1_slider_1"
-    slider_b: str = "joint_right_arm_1_hand_1_slider_2"
+    hand_roll = "joint_right_arm_1_hand_1_x4_1_dof_x4"
+    hand_grip = "joint_right_arm_1_hand_1_x4_2_dof_x4"
+    slider_a = "joint_right_arm_1_hand_1_slider_1"
+    slider_b = "joint_right_arm_1_hand_1_slider_2"
 
 
 class RightArm(Node):
-    shoulder_yaw: str = "joint_right_arm_1_x8_1_dof_x8"
-    shoulder_pitch: str = "joint_right_arm_1_x8_2_dof_x8"
-    shoulder_roll: str = "joint_right_arm_1_x6_1_dof_x6"
-    elbow_yaw: str = "joint_right_arm_1_x6_2_dof_x6"
-    elbow_roll: str = "joint_right_arm_1_x4_1_dof_x4"
-    hand: Node = RightHand()
+    shoulder_yaw = "joint_right_arm_1_x8_1_dof_x8"
+    shoulder_pitch = "joint_right_arm_1_x8_2_dof_x8"
+    shoulder_roll = "joint_right_arm_1_x6_1_dof_x6"
+    elbow_yaw = "joint_right_arm_1_x6_2_dof_x6"
+    elbow_roll = "joint_right_arm_1_x4_1_dof_x4"
+    hand = RightHand()
 
 
 class LeftLeg(Node):
-    hip_roll: str = "joint_legs_1_x8_2_dof_x8"
-    hip_yaw: str = "joint_legs_1_left_leg_1_x8_1_dof_x8"
-    hip_pitch: str = "joint_legs_1_left_leg_1_x10_1_dof_x10"
-    knee_motor: str = "joint_legs_1_left_leg_1_x10_2_dof_x10"
-    knee: str = "joint_legs_1_left_leg_1_knee_revolute"
-    ankle_motor: str = "joint_legs_1_left_leg_1_x6_1_dof_x6"
-    ankle: str = "joint_legs_1_left_leg_1_ankle_revolute"
-    foot_roll: str = "joint_legs_1_left_leg_1_x4_1_dof_x4"
+    hip_roll = "joint_legs_1_x8_2_dof_x8"
+    hip_yaw = "joint_legs_1_left_leg_1_x8_1_dof_x8"
+    hip_pitch = "joint_legs_1_left_leg_1_x10_1_dof_x10"
+    knee_motor = "joint_legs_1_left_leg_1_x10_2_dof_x10"
+    knee = "joint_legs_1_left_leg_1_knee_revolute"
+    ankle_motor = "joint_legs_1_left_leg_1_x6_1_dof_x6"
+    ankle = "joint_legs_1_left_leg_1_ankle_revolute"
+    foot_roll = "joint_legs_1_left_leg_1_x4_1_dof_x4"
 
 
 class RightLeg(Node):
-    hip_roll: str = "joint_legs_1_x8_1_dof_x8"
-    hip_yaw: str = "joint_legs_1_right_leg_1_x8_1_dof_x8"
-    hip_pitch: str = "joint_legs_1_right_leg_1_x10_2_dof_x10"
-    knee_motor: str = "joint_legs_1_right_leg_1_x10_1_dof_x10"
-    knee: str = "joint_legs_1_right_leg_1_knee_revolute"
-    ankle_motor: str = "joint_legs_1_right_leg_1_x6_1_dof_x6"
-    ankle: str = "joint_legs_1_right_leg_1_ankle_revolute"
-    foot_roll: str = "joint_legs_1_right_leg_1_x4_1_dof_x4"
+    hip_roll = "joint_legs_1_x8_1_dof_x8"
+    hip_yaw = "joint_legs_1_right_leg_1_x8_1_dof_x8"
+    hip_pitch = "joint_legs_1_right_leg_1_x10_2_dof_x10"
+    knee_motor = "joint_legs_1_right_leg_1_x10_1_dof_x10"
+    knee = "joint_legs_1_right_leg_1_knee_revolute"
+    ankle_motor = "joint_legs_1_right_leg_1_x6_1_dof_x6"
+    ankle = "joint_legs_1_right_leg_1_ankle_revolute"
+    foot_roll = "joint_legs_1_right_leg_1_x4_1_dof_x4"
 
 
 class Legs(Node):
-    left: Node = LeftLeg()
-    right: Node = RightLeg()
+    left = LeftLeg()
+    right = RightLeg()
 
 
 class Stompy(Node):
-    head: Node = Head()
-    torso: Node = Torso()
-    left_arm: Node = LeftArm()
-    right_arm: Node = RightArm()
-    legs: Node = Legs()
+    head = Head()
+    torso = Torso()
+    left_arm = LeftArm()
+    right_arm = RightArm()
+    legs = Legs()
+
+    @classmethod
+    def default_positions(cls) -> Dict[str, float]:
+        return {
+            Stompy.head.left_right: np.deg2rad(-54),
+            Stompy.head.up_down: 0.0,
+            Stompy.torso.pitch: 0.0,
+            Stompy.left_arm.shoulder_yaw: np.deg2rad(60),
+            Stompy.left_arm.shoulder_pitch: np.deg2rad(60),
+            Stompy.right_arm.shoulder_yaw: np.deg2rad(-60),
+        }
 
 
 def print_joints() -> None: