From 5201a177b5b2a550c50008315ded1e4cec681759 Mon Sep 17 00:00:00 2001 From: Simon Le Cleach Date: Tue, 6 Feb 2024 17:43:15 -0500 Subject: [PATCH 1/2] example implemented --- examples/assets/rotor.xml | 45 +++++++++++++ examples/env/env_rotor.py | 80 +++++++++++++++++++++++ examples/env/environment.py | 125 ++++++++++++++++++------------------ 3 files changed, 188 insertions(+), 62 deletions(-) create mode 100644 examples/assets/rotor.xml create mode 100644 examples/env/env_rotor.py diff --git a/examples/assets/rotor.xml b/examples/assets/rotor.xml new file mode 100644 index 0000000000..9a31c940ea --- /dev/null +++ b/examples/assets/rotor.xml @@ -0,0 +1,45 @@ + + + + + + diff --git a/examples/env/env_rotor.py b/examples/env/env_rotor.py new file mode 100644 index 0000000000..9c00b401c4 --- /dev/null +++ b/examples/env/env_rotor.py @@ -0,0 +1,80 @@ +#%% +# Copyright (c) 2022 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. + +########################################################################### +# Rotor environment +########################################################################### + +import os +import math +import warp as wp +import warp.sim +import matplotlib.pyplot as plt +import numpy as np + +from environment import Environment, run_env, RenderMode, IntegratorType + + +class RotorEnvironment(Environment): + sim_name = "env_rotor" + env_offset = (2.0, 0.0, 2.0) + opengl_render_settings = dict(scaling=1.0) + usd_render_settings = dict(scaling=100.0) + + sim_substeps_euler = 32 + sim_substeps_xpbd = 5 + + xpbd_settings = dict( + iterations=2, + joint_linear_relaxation=0.7, + joint_angular_relaxation=0.5, + rigid_contact_relaxation=1.0, + rigid_contact_con_weighting=True, + ) + + def create_articulation(self, builder): + wp.sim.parse_mjcf( + os.path.join(os.path.dirname(__file__), "../assets/rotor.xml"), + builder, + stiffness=0.0, + damping=0.1, + armature=0.007, + armature_scale=10.0, + contact_ke=1.0e4, + contact_kd=1.0e2, + contact_kf=1.0e2, + contact_mu=0.5, + contact_restitution=0.0, + limit_ke=1.0e2, + limit_kd=1.0e1, + enable_self_collisions=True, + up_axis="y", + ) + + builder.joint_q[:] = [0.0] + builder.joint_qd[:] = [-1.0] + # builder.joint_target[:] = [0.0] + +#%% +env = RotorEnvironment() +env.integrator_type = IntegratorType.XPBD +env.render_mode = RenderMode.NONE +env.num_envs = 1 +env.profile = False +env.plot_joint_coords = True +env.init() + +#%% +plot_joint_coords = env.run() +plt.plot(plot_joint_coords) + +#%% +# if __name__ == "__main__": + # run_env(RotorEnvironment) + +# %% diff --git a/examples/env/environment.py b/examples/env/environment.py index c2f158f600..fa04b8e55d 100644 --- a/examples/env/environment.py +++ b/examples/env/environment.py @@ -425,68 +425,69 @@ def run(self): ax[i, 6].yaxis.get_major_locator().set_params(integer=True) plt.show() - if self.plot_joint_coords: - import matplotlib.pyplot as plt - - joint_q_history = np.array(joint_q_history) - dof_q = joint_q_history.shape[1] - ncols = int(np.ceil(np.sqrt(dof_q))) - nrows = int(np.ceil(dof_q / float(ncols))) - fig, axes = plt.subplots( - ncols=ncols, - nrows=nrows, - constrained_layout=True, - figsize=(ncols * 3.5, nrows * 3.5), - squeeze=False, - sharex=True, - ) - - joint_id = 0 - joint_type_names = { - wp.sim.JOINT_BALL: "ball", - wp.sim.JOINT_REVOLUTE: "hinge", - wp.sim.JOINT_PRISMATIC: "slide", - wp.sim.JOINT_UNIVERSAL: "universal", - wp.sim.JOINT_COMPOUND: "compound", - wp.sim.JOINT_FREE: "free", - wp.sim.JOINT_FIXED: "fixed", - wp.sim.JOINT_DISTANCE: "distance", - wp.sim.JOINT_D6: "D6", - } - joint_lower = self.model.joint_limit_lower.numpy() - joint_upper = self.model.joint_limit_upper.numpy() - joint_type = self.model.joint_type.numpy() - while joint_id < len(joint_type) - 1 and joint_type[joint_id] == wp.sim.JOINT_FIXED: - # skip fixed joints - joint_id += 1 - q_start = self.model.joint_q_start.numpy() - qd_start = self.model.joint_qd_start.numpy() - qd_i = qd_start[joint_id] - for dim in range(ncols * nrows): - ax = axes[dim // ncols, dim % ncols] - if dim >= dof_q: - ax.axis("off") - continue - ax.grid() - ax.plot(joint_q_history[:, dim]) - if joint_type[joint_id] != wp.sim.JOINT_FREE: - lower = joint_lower[qd_i] - if abs(lower) < 2 * np.pi: - ax.axhline(lower, color="red") - upper = joint_upper[qd_i] - if abs(upper) < 2 * np.pi: - ax.axhline(upper, color="red") - joint_name = joint_type_names[joint_type[joint_id]] - ax.set_title(f"$\\mathbf{{q_{{{dim}}}}}$ ({self.model.joint_name[joint_id]} / {joint_name} {joint_id})") - if joint_id < self.model.joint_count - 1 and q_start[joint_id + 1] == dim + 1: - joint_id += 1 - qd_i = qd_start[joint_id] - else: - qd_i += 1 - plt.tight_layout() - plt.show() - - return 1000.0 * float(self.num_envs) / avg_time + # if self.plot_joint_coords: + # import matplotlib.pyplot as plt + + # joint_q_history = np.array(joint_q_history) + # dof_q = joint_q_history.shape[1] + # ncols = int(np.ceil(np.sqrt(dof_q))) + # nrows = int(np.ceil(dof_q / float(ncols))) + # fig, axes = plt.subplots( + # ncols=ncols, + # nrows=nrows, + # constrained_layout=True, + # figsize=(ncols * 3.5, nrows * 3.5), + # squeeze=False, + # sharex=True, + # ) + + # joint_id = 0 + # joint_type_names = { + # wp.sim.JOINT_BALL: "ball", + # wp.sim.JOINT_REVOLUTE: "hinge", + # wp.sim.JOINT_PRISMATIC: "slide", + # wp.sim.JOINT_UNIVERSAL: "universal", + # wp.sim.JOINT_COMPOUND: "compound", + # wp.sim.JOINT_FREE: "free", + # wp.sim.JOINT_FIXED: "fixed", + # wp.sim.JOINT_DISTANCE: "distance", + # wp.sim.JOINT_D6: "D6", + # } + # joint_lower = self.model.joint_limit_lower.numpy() + # joint_upper = self.model.joint_limit_upper.numpy() + # joint_type = self.model.joint_type.numpy() + # while joint_id < len(joint_type) - 1 and joint_type[joint_id] == wp.sim.JOINT_FIXED: + # # skip fixed joints + # joint_id += 1 + # q_start = self.model.joint_q_start.numpy() + # qd_start = self.model.joint_qd_start.numpy() + # qd_i = qd_start[joint_id] + # for dim in range(ncols * nrows): + # ax = axes[dim // ncols, dim % ncols] + # if dim >= dof_q: + # ax.axis("off") + # continue + # ax.grid() + # ax.plot(joint_q_history[:, dim]) + # if joint_type[joint_id] != wp.sim.JOINT_FREE: + # lower = joint_lower[qd_i] + # if abs(lower) < 2 * np.pi: + # ax.axhline(lower, color="red") + # upper = joint_upper[qd_i] + # if abs(upper) < 2 * np.pi: + # ax.axhline(upper, color="red") + # joint_name = joint_type_names[joint_type[joint_id]] + # ax.set_title(f"$\\mathbf{{q_{{{dim}}}}}$ ({self.model.joint_name[joint_id]} / {joint_name} {joint_id})") + # if joint_id < self.model.joint_count - 1 and q_start[joint_id + 1] == dim + 1: + # joint_id += 1 + # qd_i = qd_start[joint_id] + # else: + # qd_i += 1 + # plt.tight_layout() + # plt.show() + + # return 1000.0 * float(self.num_envs) / avg_time + return joint_q_history def run_env(Demo): From fdd5ef0d84c766735bb7d416e2e3e904c237b1d0 Mon Sep 17 00:00:00 2001 From: Simon Le Cleach Date: Tue, 6 Feb 2024 17:47:29 -0500 Subject: [PATCH 2/2] example implemented --- examples/env/env_rotor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/env/env_rotor.py b/examples/env/env_rotor.py index 9c00b401c4..46bb573228 100644 --- a/examples/env/env_rotor.py +++ b/examples/env/env_rotor.py @@ -57,8 +57,7 @@ def create_articulation(self, builder): ) builder.joint_q[:] = [0.0] - builder.joint_qd[:] = [-1.0] - # builder.joint_target[:] = [0.0] + builder.joint_qd[:] = [-2.0] #%% env = RotorEnvironment() @@ -72,7 +71,8 @@ def create_articulation(self, builder): #%% plot_joint_coords = env.run() plt.plot(plot_joint_coords) - +plt.xlabel("time_step") +plt.ylabel("angle in radian") #%% # if __name__ == "__main__": # run_env(RotorEnvironment)