Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Devel #55

Merged
merged 10 commits into from
Oct 16, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
feat: add network file and adapt to anymal_d
  • Loading branch information
mqjinwon committed Oct 15, 2024
commit 6311c204f6d04319dc6232a5eec86d98c1f78f37
158 changes: 63 additions & 95 deletions exts/StrideSim/StrideSim/anymal_articulation.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,3 @@
# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#

import io
import numpy as np
import sys
Expand All @@ -20,8 +11,8 @@
from omni.isaac.core.utils.prims import define_prim, get_prim_at_path
from omni.isaac.core.utils.rotations import euler_to_rot_matrix, quat_to_euler_angles, quat_to_rot_matrix
from omni.isaac.core.utils.stage import get_current_stage
from omni.isaac.core.utils.types import ArticulationAction
from omni.isaac.nucleus import get_assets_root_path
from omni.isaac.quadruped.utils import LstmSeaNetwork

from StrideSim.settings import ISAACLAB_LAB, ISAACLAB_LAB_ASSETS, ISAACLAB_LAB_TASKS

Expand All @@ -35,21 +26,19 @@
from omni.isaac.lab.utils.assets import ISAACLAB_NUCLEUS_DIR


class AnymalD_Atriculation(Articulation):
"""The ANYmal quadruped"""
class AnymalD_Atriculation:
"""The AnymalD quadruped"""

def __init__(
self,
prim_path: str,
name: str = "anymal",
name: str = "anymalD",
usd_path: Optional[str] = None,
position: Optional[np.ndarray] = None,
orientation: Optional[np.ndarray] = None,
) -> None:
"""
[Summary]

initialize robot, set up sensors and controller
Initialize robot and load RL policy.

Args:
prim_path {str} -- prim path of the robot on the stage
Expand All @@ -74,37 +63,36 @@ def __init__(

asset_path = f"{ISAACLAB_NUCLEUS_DIR}/Robots/ANYbotics/ANYmal-D/anymal_d.usd"

carb.log_warn("asset path is: " + asset_path)
prim.GetReferences().AddReference(asset_path)

super().__init__(prim_path=self._prim_path, name=name, position=position, orientation=orientation)
self.robot = Articulation(
prim_path=self._prim_path,
name=name,
position=position,
orientation=orientation,
)

self._dof_control_modes: List[int] = list()

# print("assets_root_path: ", assets_root_path)

# Policy
file_content = omni.client.read_file(assets_root_path + "/Isaac/Samples/Quadruped/Anymal_Policies/policy_1.pt")[
2
]
file_content = omni.client.read_file(
"/home/jin/Documents/StrideSim/exts/StrideSim/StrideSim/network/policy.pt"
)[2]
file = io.BytesIO(memoryview(file_content).tobytes())

self._policy = torch.jit.load(file)
self._base_vel_lin_scale = 2.0
self._base_vel_ang_scale = 0.25
self._joint_pos_scale = 1.0
self._joint_vel_scale = 0.05
self._base_vel_lin_scale = 1
self._base_vel_ang_scale = 1
self._action_scale = 0.5
self._default_joint_pos = np.array([0.0, 0.4, -0.8, 0.0, -0.4, 0.8, -0.0, 0.4, -0.8, -0.0, -0.4, 0.8])
self._default_joint_pos = np.array(
# [0.0, 0.4, -0.8, 0.0, -0.4, 0.8, -0.0, 0.4, -0.8, -0.0, -0.4, 0.8]
[0.0, 0.0, 0.0, 0.0, 0.4, -0.4, 0.4, -0.4, -0.8, 0.8, -0.8, 0.8]
)
self._previous_action = np.zeros(12)
self._policy_counter = 0

# Actuator network
file_content = omni.client.read_file(
assets_root_path + "/Isaac/Samples/Quadruped/Anymal_Policies/sea_net_jit2.pt"
)[2]
file = io.BytesIO(memoryview(file_content).tobytes())
self._actuator_network = LstmSeaNetwork()
self._actuator_network.setup(file, self._default_joint_pos)
self._actuator_network.reset()
self._decimation = 4

# Height scaner
y = np.arange(-0.5, 0.6, 0.1)
Expand All @@ -123,9 +111,8 @@ def _hit_report_callback(self, hit):
return True

def _compute_observation(self, command):
"""[summary]

compute the observation vector for the policy
"""
Compute the observation vector for the policy

Argument:
command {np.ndarray} -- the robot command (v_x, v_y, w_z)
Expand All @@ -134,12 +121,20 @@ def _compute_observation(self, command):
np.ndarray -- The observation vector.

"""
lin_vel_I = self.get_linear_velocity()
ang_vel_I = self.get_angular_velocity()
pos_IB, q_IB = self.get_world_pose()
lin_vel_I = self.robot.get_linear_velocity()
ang_vel_I = self.robot.get_angular_velocity()
pos_IB, q_IB = self.robot.get_world_pose()

R_IB = quat_to_rot_matrix(q_IB)
R_BI = R_IB.transpose()

print("lin_vel_I: ", lin_vel_I)
print("ang_vel_I: ", ang_vel_I)
print("pos_IB: ", pos_IB)
print("q_IB: ", q_IB)
print("R_IB: ", R_IB)
print("R_BI: ", R_BI)

lin_vel_b = np.matmul(R_BI, lin_vel_I)
ang_vel_b = np.matmul(R_BI, ang_vel_I)
gravity_b = np.matmul(R_BI, np.array([0.0, 0.0, -1.0]))
Expand All @@ -156,24 +151,12 @@ def _compute_observation(self, command):
obs[10] = self._base_vel_lin_scale * command[1]
obs[11] = self._base_vel_ang_scale * command[2]
# Joint states
# joint_state from the DC interface now has the order of
# 'FL_hip_joint', 'FR_hip_joint', 'RL_hip_joint', 'RR_hip_joint',
# 'FL_thigh_joint', 'FR_thigh_joint', 'RL_thigh_joint', 'RR_thigh_joint',
# 'FL_calf_joint', 'FR_calf_joint', 'RL_calf_joint', 'RR_calf_joint'

# while the learning controller uses the order of
# FL_hip_joint FL_thigh_joint FL_calf_joint
# FR_hip_joint FR_thigh_joint FR_calf_joint
# RL_hip_joint RL_thigh_joint RL_calf_joint
# RR_hip_joint RR_thigh_joint RR_calf_joint
# Convert DC order to controller order for joint info
current_joint_pos = self.get_joint_positions()
current_joint_vel = self.get_joint_velocities()
current_joint_pos = np.array(current_joint_pos.reshape([3, 4]).T.flat)
current_joint_vel = np.array(current_joint_vel.reshape([3, 4]).T.flat)
obs[12:24] = self._joint_pos_scale * (current_joint_pos - self._default_joint_pos)
obs[24:36] = self._joint_vel_scale * current_joint_vel
current_joint_pos = self.robot.get_joint_positions()
current_joint_vel = self.robot.get_joint_velocities()
obs[12:24] = current_joint_pos - self._default_joint_pos
obs[24:36] = current_joint_vel

# Previous Action
obs[36:48] = self._previous_action

# height_scanner
Expand All @@ -186,69 +169,54 @@ def _compute_observation(self, command):
for i in range(world_scan_points.shape[0]):
self._query_info.clear()
self.physx_query_interface.raycast_all(
tuple(world_scan_points[i]), (0.0, 0.0, -1.0), 100, self._hit_report_callback
tuple(world_scan_points[i]),
(0.0, 0.0, -1.0),
100,
self._hit_report_callback,
)
if self._query_info:
distance = min(self._query_info)
obs[48 + i] = np.clip(distance - 0.5, -1.0, 1.0)
else:
print("No hit")
pass
# print("No hit")
return obs

def advance(self, dt, command):
"""[summary]

compute the desired torques and apply them to the articulation
"""
Compute the desired torques and apply them to the articulation

Argument:
dt {float} -- Timestep update in the world.
command {np.ndarray} -- the robot command (v_x, v_y, w_z)

"""
if self._policy_counter % 4 == 0:
if self._policy_counter % self._decimation == 0:
obs = self._compute_observation(command)
with torch.no_grad():
obs = torch.from_numpy(obs).view(1, -1).float()
self.action = self._policy(obs).detach().view(-1).numpy()
self._previous_action = self.action.copy()

# joint_state from the DC interface now has the order of
# 'FL_hip_joint', 'FR_hip_joint', 'RL_hip_joint', 'RR_hip_joint',
# 'FL_thigh_joint', 'FR_thigh_joint', 'RL_thigh_joint', 'RR_thigh_joint',
# 'FL_calf_joint', 'FR_calf_joint', 'RL_calf_joint', 'RR_calf_joint'

# while the learning controller uses the order of
# FL_hip_joint FL_thigh_joint FL_calf_joint
# FR_hip_joint FR_thigh_joint FR_calf_joint
# RL_hip_joint RL_thigh_joint RL_calf_joint
# RR_hip_joint RR_thigh_joint RR_calf_joint
# Convert DC order to controller order for joint info
current_joint_pos = self.get_joint_positions()
current_joint_vel = self.get_joint_velocities()
current_joint_pos = np.array(current_joint_pos.reshape([3, 4]).T.flat)
current_joint_vel = np.array(current_joint_vel.reshape([3, 4]).T.flat)
joint_torques, _ = self._actuator_network.compute_torques(
current_joint_pos, current_joint_vel, self._action_scale * self.action
)
action = ArticulationAction(joint_positions=self._default_joint_pos + (self.action * self._action_scale))

# finally convert controller order to DC order for command torque
torque_reorder = np.array(joint_torques.reshape([4, 3]).T.flat)
self.set_joint_efforts(torque_reorder)
print("network action: ", self.action)

self.robot.apply_action(action)

self._policy_counter += 1

def initialize(self, physics_sim_view=None) -> None:
"""[summary]

initialize the articulation interface, set up drive mode
"""
super().initialize(physics_sim_view=physics_sim_view)
self.get_articulation_controller().set_effort_modes("force")
self.get_articulation_controller().switch_control_mode("effort")
Initialize robot the articulation interface, set up drive mode
"""
self.robot.initialize(physics_sim_view=physics_sim_view)
self.robot.get_articulation_controller().set_effort_modes("force")
self.robot.get_articulation_controller().switch_control_mode("position")
self.robot._articulation_view.set_gains(np.zeros(12) + 200, np.zeros(12) + 5)

def post_reset(self) -> None:
"""[summary]

post reset articulation
"""
super().post_reset()
Post reset articulation
"""
self.robot.post_reset()
Loading
Loading