From db7230c3ec2679701c0c36bc68872f48dd8bde80 Mon Sep 17 00:00:00 2001 From: Neko Asakura Date: Fri, 20 Dec 2024 04:13:27 +1000 Subject: [PATCH 01/35] ext/pyrender/viewer: clean up unused imports --- genesis/ext/pyrender/viewer.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/genesis/ext/pyrender/viewer.py b/genesis/ext/pyrender/viewer.py index f9f68cc5..4a08d36b 100644 --- a/genesis/ext/pyrender/viewer.py +++ b/genesis/ext/pyrender/viewer.py @@ -14,15 +14,8 @@ import genesis as gs -try: - from Tkinter import Tk - from Tkinter import tkFileDialog as filedialog -except Exception: - try: - from tkinter import Tk - from tkinter import filedialog as filedialog - except Exception: - pass +from tkinter import Tk +from tkinter import filedialog import pyglet from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter From d8f6bedf922d0f1fa4f2a408ed7e21c02c7a55a8 Mon Sep 17 00:00:00 2001 From: Neko Asakura Date: Fri, 20 Dec 2024 04:18:37 +1000 Subject: [PATCH 02/35] ext/pyrender/viewer: resolve macOS crashes when initialising Tk() --- genesis/ext/pyrender/viewer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/genesis/ext/pyrender/viewer.py b/genesis/ext/pyrender/viewer.py index 4a08d36b..c2f2e2da 100644 --- a/genesis/ext/pyrender/viewer.py +++ b/genesis/ext/pyrender/viewer.py @@ -16,6 +16,8 @@ from tkinter import Tk from tkinter import filedialog +root = Tk() +root.withdraw() import pyglet from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter @@ -961,7 +963,6 @@ def _get_save_filename(self, file_exts): } filetypes = [file_types[x] for x in file_exts] try: - root = Tk() save_dir = self.viewer_flags["save_directory"] if save_dir is None: save_dir = os.getcwd() @@ -971,7 +972,6 @@ def _get_save_filename(self, file_exts): except Exception: return None - root.destroy() if filename == (): return None return filename From 6f9dcc67cf793f8526676aac4e97de3ff929a4c8 Mon Sep 17 00:00:00 2001 From: Johnson Wang Date: Sun, 22 Dec 2024 06:48:29 -0500 Subject: [PATCH 03/35] allow to toggle batching link/dof info --- .../entities/rigid_entity/rigid_entity.py | 18 +- .../entities/rigid_entity/rigid_joint.py | 3 +- .../engine/solvers/rigid/collider_decomp.py | 23 +- .../solvers/rigid/constraint_solver_decomp.py | 23 +- .../rigid/constraint_solver_decomp_island.py | 26 +- .../engine/solvers/rigid/contact_island.py | 13 +- .../solvers/rigid/rigid_solver_decomp.py | 369 +++++++++++------- genesis/options/solvers.py | 4 + 8 files changed, 298 insertions(+), 181 deletions(-) diff --git a/genesis/engine/entities/rigid_entity/rigid_entity.py b/genesis/engine/entities/rigid_entity/rigid_entity.py index 0517c80a..ce53344a 100644 --- a/genesis/engine/entities/rigid_entity/rigid_entity.py +++ b/genesis/engine/entities/rigid_entity/rigid_entity.py @@ -666,7 +666,8 @@ def _func_get_jacobian(self, tgt_link_idx, i_b, pos_mask, rot_mask): tgt_link_pos = self._solver.links_state[tgt_link_idx, i_b].pos i_l = tgt_link_idx while i_l > -1: - l_info = self._solver.links_info[i_l] + I_l = [i_l, i_b] if ti.static(self.solver._options.batch_links_info) else i_l + l_info = self._solver.links_info[I_l] l_state = self._solver.links_state[i_l, i_b] if l_info.joint_type == gs.JOINT_TYPE.FIXED: @@ -674,8 +675,9 @@ def _func_get_jacobian(self, tgt_link_idx, i_b, pos_mask, rot_mask): elif l_info.joint_type == gs.JOINT_TYPE.REVOLUTE: i_d = l_info.dof_start + I_d = [i_d, i_b] if ti.static(self.solver._options.batch_dofs_info) else i_d i_d_jac = i_d - self._dof_start - rotation = gu.ti_transform_by_quat(self._solver.dofs_info[i_d].motion_ang, l_state.quat) + rotation = gu.ti_transform_by_quat(self._solver.dofs_info[I_d].motion_ang, l_state.quat) translation = rotation.cross(tgt_link_pos - l_state.pos) self._jacobian[0, i_d_jac, i_b] = translation[0] * pos_mask[0] @@ -687,8 +689,9 @@ def _func_get_jacobian(self, tgt_link_idx, i_b, pos_mask, rot_mask): elif l_info.joint_type == gs.JOINT_TYPE.PRISMATIC: i_d = l_info.dof_start + I_d = [i_d, i_b] if ti.static(self.solver._options.batch_dofs_info) else i_d i_d_jac = i_d - self._dof_start - translation = gu.ti_transform_by_quat(self._solver.dofs_info[i_d].motion_vel, l_state.quat) + translation = gu.ti_transform_by_quat(self._solver.dofs_info[I_d].motion_vel, l_state.quat) self._jacobian[0, i_d_jac, i_b] = translation[0] * pos_mask[0] self._jacobian[1, i_d_jac, i_b] = translation[1] * pos_mask[1] @@ -706,7 +709,8 @@ def _func_get_jacobian(self, tgt_link_idx, i_b, pos_mask, rot_mask): for i_d_ in range(3): i_d = l_info.dof_start + i_d_ + 3 i_d_jac = i_d - self._dof_start - rotation = self._solver.dofs_info[i_d].motion_ang + I_d = [i_d, i_b] if ti.static(self.solver._options.batch_dofs_info) else i_d + rotation = self._solver.dofs_info[I_d].motion_ang translation = rotation.cross(tgt_link_pos - l_state.pos) self._jacobian[0, i_d_jac, i_b] = translation[0] * pos_mask[0] @@ -1156,8 +1160,10 @@ def _kernel_inverse_kinematics( # Resample init q if respect_joint_limit and i_sample < max_samples - 1: for i_l in range(self.link_start, self.link_end): - l_info = self._solver.links_info[i_l] - dof_info = self._solver.dofs_info[l_info.dof_start] + I_l = [i_l, i_b] if ti.static(self.solver._options.batch_links_info) else i_l + l_info = self._solver.links_info[I_l] + I_dof_start = [l_info.dof_start, i_b] if ti.static(self.solver._options.batch_dofs_info) else l_info.dof_start + dof_info = self._solver.dofs_info[I_dof_start] q_start = l_info.q_start if l_info.joint_type == gs.JOINT_TYPE.FREE: diff --git a/genesis/engine/entities/rigid_entity/rigid_joint.py b/genesis/engine/entities/rigid_entity/rigid_joint.py index 660971d9..6e57c4a0 100644 --- a/genesis/engine/entities/rigid_entity/rigid_joint.py +++ b/genesis/engine/entities/rigid_entity/rigid_joint.py @@ -90,7 +90,8 @@ def get_pos(self): def _kernel_get_pos(self, tensor: ti.types.ndarray()): for i_b in range(self._solver._B): - l_info = self._solver.links_info[self._idx, i_b] + I_l = [self._idx, i_b] if ti.static(self._solver._options.batch_links_info) else self._idx + l_info = self._solver.links_info[I_l] i_p = l_info.parent_idx p_pos = ti.Vector.zero(gs.ti_float, 3) diff --git a/genesis/engine/solvers/rigid/collider_decomp.py b/genesis/engine/solvers/rigid/collider_decomp.py index 968ad214..6b438750 100644 --- a/genesis/engine/solvers/rigid/collider_decomp.py +++ b/genesis/engine/solvers/rigid/collider_decomp.py @@ -164,10 +164,13 @@ def clear(self): i_la = self.contact_data[i_c, i_b].link_a i_lb = self.contact_data[i_c, i_b].link_b + I_la = [i_la, i_b] if ti.static(self._solver._options.batch_links_info) else i_la + I_lb = [i_lb, i_b] if ti.static(self._solver._options.batch_links_info) else i_lb + # pair of hibernated-fixed links -> hibernated contact # TODO: we should also include hibernated-hibernated links and wake up the whole contact island once a new collision is detected - if (self._solver.links_state[i_la, i_b].hibernated and self._solver.links_info[i_lb].is_fixed) or ( - self._solver.links_state[i_lb, i_b].hibernated and self._solver.links_info[i_la].is_fixed + if (self._solver.links_state[i_la, i_b].hibernated and self._solver.links_info[I_lb].is_fixed) or ( + self._solver.links_state[i_lb, i_b].hibernated and self._solver.links_info[I_la].is_fixed ): i_c_hibernated = self.n_contacts_hibernated[i_b] if i_c != i_c_hibernated: @@ -549,6 +552,8 @@ def _func_update_aabbs(self): def _func_check_collision_valid(self, i_ga, i_gb, i_b): i_la = self._solver.geoms_info[i_ga].link_idx i_lb = self._solver.geoms_info[i_gb].link_idx + I_la = [i_la, i_b] if ti.static(self._solver._options.batch_links_info) else i_la + I_lb = [i_lb, i_b] if ti.static(self._solver._options.batch_links_info) else i_lb is_valid = True # geoms in the same link @@ -558,22 +563,22 @@ def _func_check_collision_valid(self, i_ga, i_gb, i_b): # self collision if ( ti.static(not self._solver._enable_self_collision) - and self._solver.links_info[i_la].root_idx == self._solver.links_info[i_lb].root_idx + and self._solver.links_info[I_la].root_idx == self._solver.links_info[I_lb].root_idx ): is_valid = False # adjacent links - if self._solver.links_info[i_la].parent_idx == i_lb or self._solver.links_info[i_lb].parent_idx == i_la: + if self._solver.links_info[I_la].parent_idx == i_lb or self._solver.links_info[I_lb].parent_idx == i_la: is_valid = False # pair of fixed links - if self._solver.links_info[i_la].is_fixed and self._solver.links_info[i_lb].is_fixed: + if self._solver.links_info[I_la].is_fixed and self._solver.links_info[I_lb].is_fixed: is_valid = False # hibernated <-> fixed links if ti.static(self._solver._use_hibernation): - if (self._solver.links_state[i_la, i_b].hibernated and self._solver.links_info[i_lb].is_fixed) or ( - self._solver.links_state[i_lb, i_b].hibernated and self._solver.links_info[i_la].is_fixed + if (self._solver.links_state[i_la, i_b].hibernated and self._solver.links_info[I_lb].is_fixed) or ( + self._solver.links_state[i_lb, i_b].hibernated and self._solver.links_info[I_la].is_fixed ): is_valid = False @@ -994,7 +999,9 @@ def _func_mpr(self, i_ga, i_gb, i_b): i_la = self._solver.geoms_info[i_ga].link_idx i_lb = self._solver.geoms_info[i_gb].link_idx - is_self_pair = self._solver.links_info.root_idx[i_la] == self._solver.links_info.root_idx[i_lb] + I_la = [i_la, i_b] if ti.static(self._solver._options.batch_links_info) else i_la + I_lb = [i_lb, i_b] if ti.static(self._solver._options.batch_links_info) else i_lb + is_self_pair = self._solver.links_info.root_idx[I_la] == self._solver.links_info.root_idx[I_lb] multi_contact = ( self._solver.geoms_info[i_ga].type != gs.GEOM_TYPE.SPHERE and self._solver.geoms_info[i_gb].type != gs.GEOM_TYPE.SPHERE diff --git a/genesis/engine/solvers/rigid/constraint_solver_decomp.py b/genesis/engine/solvers/rigid/constraint_solver_decomp.py index c680c31f..43c0cab1 100644 --- a/genesis/engine/solvers/rigid/constraint_solver_decomp.py +++ b/genesis/engine/solvers/rigid/constraint_solver_decomp.py @@ -105,12 +105,14 @@ def add_collision_constraints(self): impact = self._collider.contact_data[i_col, i_b] link_a = impact.link_a link_b = impact.link_b + link_a_maybe_batch = [link_a, i_b] if ti.static(self._solver._options.batch_links_info) else link_a + link_b_maybe_batch = [link_b, i_b] if ti.static(self._solver._options.batch_links_info) else link_b f = impact.friction pos = impact.pos d1, d2 = gu.orthogonals(impact.normal) - t = self._solver.links_info[link_a].invweight + self._solver.links_info[link_b].invweight * ( + t = self._solver.links_info[link_a_maybe_batch].invweight + self._solver.links_info[link_b_maybe_batch].invweight * ( link_b > -1 ) for i in range(4): @@ -143,10 +145,11 @@ def add_collision_constraints(self): link = link_b while link > -1: + link_maybe_batch = [link, i_b] if ti.static(self._solver._options.batch_links_info) else link # reverse order to make sure dofs in each row of self.jac_relevant_dofs is strictly descending - for i_d_ in range(self._solver.links_info[link].n_dofs): - i_d = self._solver.links_info[link].dof_end - 1 - i_d_ + for i_d_ in range(self._solver.links_info[link_maybe_batch].n_dofs): + i_d = self._solver.links_info[link_maybe_batch].dof_end - 1 - i_d_ cdof_ang = self._solver.dofs_state[i_d, i_b].cdof_ang cdot_vel = self._solver.dofs_state[i_d, i_b].cdof_vel @@ -164,7 +167,7 @@ def add_collision_constraints(self): self.jac_relevant_dofs[n_con, con_n_relevant_dofs, i_b] = i_d con_n_relevant_dofs += 1 - link = self._solver.links_info[link].parent_idx + link = self._solver.links_info[link_maybe_batch].parent_idx if ti.static(self.sparse_solve): self.jac_n_relevant_dofs[n_con, i_b] = con_n_relevant_dofs @@ -183,21 +186,23 @@ def add_joint_limit_constraints(self): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i_b in range(self._B): for i_l in range(self._solver.n_links): - l_info = self._solver.links_info[i_l] + I_l = [i_l, i_b] if ti.static(self._solver._options.batch_links_info) else i_l + l_info = self._solver.links_info[I_l] if l_info.joint_type == gs.JOINT_TYPE.REVOLUTE or l_info.joint_type == gs.JOINT_TYPE.PRISMATIC: i_q = l_info.q_start i_d = l_info.dof_start - pos_min = self._solver.qpos[i_q, i_b] - self._solver.dofs_info[i_d].limit[0] - pos_max = self._solver.dofs_info[i_d].limit[1] - self._solver.qpos[i_q, i_b] + I_d = [i_d, i_b] if ti.static(self._solver._options.batch_dofs_info) else i_d + pos_min = self._solver.qpos[i_q, i_b] - self._solver.dofs_info[I_d].limit[0] + pos_max = self._solver.dofs_info[I_d].limit[1] - self._solver.qpos[i_q, i_b] pos = min(min(pos_min, pos_max), 0) side = ((pos_min < pos_max) * 2 - 1) * (pos < 0) jac = side jac_qvel = jac * self._solver.dofs_state[i_d, i_b].vel - imp, aref = gu.imp_aref(self._solver.dofs_info[i_d].sol_params, pos, jac_qvel) - diag = self._solver.dofs_info[i_d].invweight * (pos < 0) * (1 - imp) / (imp + gs.EPS) + imp, aref = gu.imp_aref(self._solver.dofs_info[I_d].sol_params, pos, jac_qvel) + diag = self._solver.dofs_info[I_d].invweight * (pos < 0) * (1 - imp) / (imp + gs.EPS) aref = aref * (pos < 0) if pos < 0: n_con = self.n_constraints[i_b] diff --git a/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py b/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py index e33867a0..534686fa 100644 --- a/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py +++ b/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py @@ -131,12 +131,14 @@ def add_collision_constraints(self, island, i_b): impact = self._collider.contact_data[i_col, i_b] link_a = impact.link_a link_b = impact.link_b + link_a_maybe_batch = [link_a, i_b] if ti.static(self._solver._options.batch_links_info) else link_a + link_b_maybe_batch = [link_b, i_b] if ti.static(self._solver._options.batch_links_info) else link_b f = impact.friction pos = impact.pos d1, d2 = gu.orthogonals(impact.normal) - t = self._solver.links_info[link_a].invweight + self._solver.links_info[link_b].invweight * (link_b > -1) + t = self._solver.links_info[link_a_maybe_batch].invweight + self._solver.links_info[link_b_maybe_batch].invweight * (link_b > -1) for i in range(4): n = -d1 * f - impact.normal @@ -170,10 +172,11 @@ def add_collision_constraints(self, island, i_b): link = link_b while link > -1: + link_maybe_batch = [link, i_b] if ti.static(self._solver._options.batch_links_info) else link # reverse order to make sure dofs in each row of self.jac_relevant_dofs is strictly descending for i_d_ in range(self._solver.links_info[link].n_dofs): - i_d = self._solver.links_info[link].dof_end - 1 - i_d_ + i_d = self._solver.links_info[link_maybe_batch].dof_end - 1 - i_d_ cdof_ang = self._solver.dofs_state[i_d, i_b].cdof_ang cdot_vel = self._solver.dofs_state[i_d, i_b].cdof_vel @@ -190,7 +193,7 @@ def add_collision_constraints(self, island, i_b): self.jac_relevant_dofs[n_con, con_n_relevant_dofs, i_b] = i_d con_n_relevant_dofs += 1 - link = self._solver.links_info[link].parent_idx + link = self._solver.links_info[link_maybe_batch].parent_idx if ti.static(self.sparse_solve): self.jac_n_relevant_dofs[n_con, i_b] = con_n_relevant_dofs @@ -207,8 +210,8 @@ def add_collision_constraints(self, island, i_b): if ti.static(self._solver._use_hibernation): # wake up entities - self._solver._func_wakeup_entity(self._solver.links_info[link_a].entity_idx, i_b) - self._solver._func_wakeup_entity(self._solver.links_info[link_b].entity_idx, i_b) + self._solver._func_wakeup_entity(self._solver.links_info[link_a_maybe_batch].entity_idx, i_b) + self._solver._func_wakeup_entity(self._solver.links_info[link_b_maybe_batch].entity_idx, i_b) @ti.func def add_joint_limit_constraints(self, island, i_b): @@ -219,22 +222,23 @@ def add_joint_limit_constraints(self, island, i_b): e_info = self.entities_info[i_e] for i_l in range(e_info.link_start, e_info.link_end): - - l_info = self._solver.links_info[i_l] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + l_info = self._solver.links_info[I_l] if l_info.joint_type == gs.JOINT_TYPE.REVOLUTE or l_info.joint_type == gs.JOINT_TYPE.PRISMATIC: i_q = l_info.q_start i_d = l_info.dof_start - pos_min = self._solver.qpos[i_q, i_b] - self._solver.dofs_info[i_d].limit[0] - pos_max = self._solver.dofs_info[i_d].limit[1] - self._solver.qpos[i_q, i_b] + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d + pos_min = self._solver.qpos[i_q, i_b] - self._solver.dofs_info[I_d].limit[0] + pos_max = self._solver.dofs_info[I_d].limit[1] - self._solver.qpos[i_q, i_b] pos = min(min(pos_min, pos_max), 0) side = ((pos_min < pos_max) * 2 - 1) * (pos < 0) jac = side jac_qvel = jac * self._solver.dofs_state[i_d, i_b].vel - imp, aref = gu.imp_aref(self._solver.dofs_info[i_d].sol_params, pos, jac_qvel) - diag = self._solver.dofs_info[i_d].invweight * (pos < 0) * (1 - imp) / (imp + gs.EPS) + imp, aref = gu.imp_aref(self._solver.dofs_info[I_d].sol_params, pos, jac_qvel) + diag = self._solver.dofs_info[I_d].invweight * (pos < 0) * (1 - imp) / (imp + gs.EPS) aref = aref * (pos < 0) if pos < 0: n_con = self.n_constraints[i_b] diff --git a/genesis/engine/solvers/rigid/contact_island.py b/genesis/engine/solvers/rigid/contact_island.py index 2914b503..97ae3b93 100644 --- a/genesis/engine/solvers/rigid/contact_island.py +++ b/genesis/engine/solvers/rigid/contact_island.py @@ -66,8 +66,11 @@ def clear(self): @ti.func def add_edge(self, link_a, link_b, i_b): - ea = self.solver.links_info[link_a].entity_idx - eb = self.solver.links_info[link_b].entity_idx + link_a_maybe_batch = [link_a, i_b] if ti.static(self.solver._options.batch_links_info) else link_a + link_b_maybe_batch = [link_b, i_b] if ti.static(self.solver._options.batch_links_info) else link_b + + ea = self.solver.links_info[link_a_maybe_batch].entity_idx + eb = self.solver.links_info[link_b_maybe_batch].entity_idx self.entity_edge[ea, i_b].n = self.entity_edge[ea, i_b].n + 1 self.entity_edge[eb, i_b].n = self.entity_edge[eb, i_b].n + 1 @@ -100,9 +103,11 @@ def postprocess_island(self): impact = self.collider.contact_data[i_col, i_b] link_a = impact.link_a link_b = impact.link_b + link_a_maybe_batch = [link_a, i_b] if ti.static(self.solver._options.batch_links_info) else link_a + link_b_maybe_batch = [link_b, i_b] if ti.static(self.solver._options.batch_links_info) else link_b - ea = self.solver.links_info[link_a].entity_idx - eb = self.solver.links_info[link_b].entity_idx + ea = self.solver.links_info[link_a_maybe_batch].entity_idx + eb = self.solver.links_info[link_b_maybe_batch].entity_idx island_a = self.entity_island[ea, i_b] island_b = self.entity_island[eb, i_b] diff --git a/genesis/engine/solvers/rigid/rigid_solver_decomp.py b/genesis/engine/solvers/rigid/rigid_solver_decomp.py index a6abeb53..985c441f 100644 --- a/genesis/engine/solvers/rigid/rigid_solver_decomp.py +++ b/genesis/engine/solvers/rigid/rigid_solver_decomp.py @@ -159,6 +159,11 @@ def _init_invweight(self): dof_end = self.links_info.dof_end.to_numpy() n_dofs = self.links_info.n_dofs.to_numpy() parent_idx = self.links_info.parent_idx.to_numpy() + if self._options.batch_links_info: + dof_start = dof_start[:, 0] + dof_end = dof_end[:, 0] + n_dofs = n_dofs[:, 0] + parent_idx = parent_idx[:, 0] offsets = self.links_state.i_pos.to_numpy()[:, 0, :] @@ -209,9 +214,9 @@ def _kernel_init_invweight( invweight: ti.types.ndarray(), ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) - for i, i_b in ti.ndrange(self.n_links, self._B): - if self.links_info[i].invweight < 0: - self.links_info[i].invweight = invweight[i] + for I in ti.grouped(self.links_info): + if self.links_info[I].invweight < 0: + self.links_info[I].invweight = invweight[I[0]] def _batch_shape(self, shape=None, first_dim=False, B=None): if B is None: @@ -300,7 +305,8 @@ def _init_dof_fields(self): ctrl_mode=gs.ti_int, hibernated=gs.ti_int, # Flag for dofs that converge into a static state (hibernation) ) - self.dofs_info = struct_dof_info.field(shape=self.n_dofs_, needs_grad=False, layout=ti.Layout.SOA) + dofs_info_shape = self._batch_shape(self.n_dofs_) if self._options.batch_dofs_info else self.n_dofs_ + self.dofs_info = struct_dof_info.field(shape=dofs_info_shape, needs_grad=False, layout=ti.Layout.SOA) self.dofs_state = struct_dof_state.field( shape=self._batch_shape(self.n_dofs_), needs_grad=False, layout=ti.Layout.SOA ) @@ -339,26 +345,28 @@ def _kernel_init_dof_fields( dofs_force_range: ti.types.ndarray(), ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) - for i, b in ti.ndrange(self.n_dofs, self._B): + for I in ti.grouped(self.dofs_info): + i = I[0] # batching (if any) will be the second dim + for j in ti.static(range(3)): - self.dofs_info[i].motion_ang[j] = dofs_motion_ang[i, j] - self.dofs_info[i].motion_vel[j] = dofs_motion_vel[i, j] + self.dofs_info[I].motion_ang[j] = dofs_motion_ang[i, j] + self.dofs_info[I].motion_vel[j] = dofs_motion_vel[i, j] for j in ti.static(range(2)): - self.dofs_info[i].limit[j] = dofs_limit[i, j] - self.dofs_info[i].force_range[j] = dofs_force_range[i, j] + self.dofs_info[I].limit[j] = dofs_limit[i, j] + self.dofs_info[I].force_range[j] = dofs_force_range[i, j] for j in ti.static(range(7)): - self.dofs_info[i].sol_params[j] = dofs_sol_params[i, j] + self.dofs_info[I].sol_params[j] = dofs_sol_params[i, j] - self.dofs_info[i].sol_params[0] = self._sol_contact_resolve_time + self.dofs_info[I].sol_params[0] = self._sol_contact_resolve_time - self.dofs_info[i].armature = dofs_armature[i] - self.dofs_info[i].invweight = dofs_invweight[i] - self.dofs_info[i].stiffness = dofs_stiffness[i] - self.dofs_info[i].damping = dofs_damping[i] - self.dofs_info[i].kp = dofs_kp[i] - self.dofs_info[i].kv = dofs_kv[i] + self.dofs_info[I].armature = dofs_armature[i] + self.dofs_info[I].invweight = dofs_invweight[i] + self.dofs_info[I].stiffness = dofs_stiffness[i] + self.dofs_info[I].damping = dofs_damping[i] + self.dofs_info[I].kp = dofs_kp[i] + self.dofs_info[I].kv = dofs_kv[i] ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i, b in ti.ndrange(self.n_dofs, self._B): @@ -442,7 +450,8 @@ def _init_link_fields(self): hibernated=gs.ti_int, ) - self.links_info = struct_link_info.field(shape=self.n_links, needs_grad=False, layout=ti.Layout.SOA) + links_info_shape = self._batch_shape(self.n_links) if self._options.batch_links_info else self.n_links + self.links_info = struct_link_info.field(shape=links_info_shape, needs_grad=False, layout=ti.Layout.SOA) self.links_state = struct_link_state.field( shape=self._batch_shape(self.n_links), needs_grad=False, layout=ti.Layout.SOA ) @@ -501,38 +510,42 @@ def _kernel_init_link_fields( links_entity_idx: ti.types.ndarray(), ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) - for i, b in ti.ndrange(self.n_links, self._B): - self.links_info[i].parent_idx = links_parent_idx[i] - self.links_info[i].root_idx = links_root_idx[i] - self.links_info[i].q_start = links_q_start[i] - self.links_info[i].dof_start = links_dof_start[i] - self.links_info[i].q_end = links_q_end[i] - self.links_info[i].dof_end = links_dof_end[i] - self.links_info[i].n_dofs = links_dof_end[i] - links_dof_start[i] - self.links_info[i].joint_type = links_joint_type[i] - self.links_info[i].invweight = links_invweight[i] - self.links_info[i].is_fixed = links_is_fixed[i] - self.links_info[i].entity_idx = links_entity_idx[i] + for I in ti.grouped(self.links_info): + i = I[0] + + self.links_info[I].parent_idx = links_parent_idx[i] + self.links_info[I].root_idx = links_root_idx[i] + self.links_info[I].q_start = links_q_start[i] + self.links_info[I].dof_start = links_dof_start[i] + self.links_info[I].q_end = links_q_end[i] + self.links_info[I].dof_end = links_dof_end[i] + self.links_info[I].n_dofs = links_dof_end[i] - links_dof_start[i] + self.links_info[I].joint_type = links_joint_type[i] + self.links_info[I].invweight = links_invweight[i] + self.links_info[I].is_fixed = links_is_fixed[i] + self.links_info[I].entity_idx = links_entity_idx[i] for j in ti.static(range(4)): - self.links_info[i].quat[j] = links_quat[i, j] - self.links_info[i].joint_quat[j] = links_joint_quat[i, j] - self.links_info[i].inertial_quat[j] = links_inertial_quat[i, j] + self.links_info[I].quat[j] = links_quat[i, j] + self.links_info[I].joint_quat[j] = links_joint_quat[i, j] + self.links_info[I].inertial_quat[j] = links_inertial_quat[i, j] for j in ti.static(range(3)): - self.links_info[i].pos[j] = links_pos[i, j] - self.links_info[i].joint_pos[j] = links_joint_pos[i, j] - self.links_info[i].inertial_pos[j] = links_inertial_pos[i, j] + self.links_info[I].pos[j] = links_pos[i, j] + self.links_info[I].joint_pos[j] = links_joint_pos[i, j] + self.links_info[I].inertial_pos[j] = links_inertial_pos[i, j] - self.links_info[i].inertial_mass = links_inertial_mass[i] + self.links_info[I].inertial_mass = links_inertial_mass[i] for j1 in ti.static(range(3)): for j2 in ti.static(range(3)): - self.links_info[i].inertial_i[j1, j2] = links_inertial_i[i, j1, j2] + self.links_info[I].inertial_i[j1, j2] = links_inertial_i[i, j1, j2] ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i, b in ti.ndrange(self.n_links, self._B): + I = [i, b] if ti.static(self._options.batch_links_info) else i + # Update state for root fixed link. Their state will not be updated in forward kinematics later but can be manually changed by user. - if self.links_info[i].parent_idx == -1 and self.links_info[i].is_fixed: + if self.links_info[I].parent_idx == -1 and self.links_info[I].is_fixed: for j in ti.static(range(4)): self.links_state[i, b].quat[j] = links_quat[i, j] @@ -836,10 +849,11 @@ def _kernel_adjust_link_inertia( link_idx: ti.i32, ratio: ti.f32, ): - self.links_info[link_idx].invweight /= ratio - self.links_info[link_idx].inertial_mass *= ratio - for j1, j2 in ti.ndrange(3, 3): - self.links_info[link_idx].inertial_i[j1, j2] *= ratio + for I_l in ti.grouped(self.links_info): + self.links_info[I_l].invweight /= ratio + self.links_info[I_l].inertial_mass *= ratio + for j1, j2 in ti.ndrange(3, 3): + self.links_info[I_l].inertial_i[j1, j2] *= ratio def _init_vgeom_fields(self): struct_vgeom_info = ti.types.struct( @@ -971,8 +985,13 @@ def _kernel_init_entity_fields( self.entities_info[i].gravity_compensation = entities_gravity_compensation[i] - for i_d in range(entities_dof_start[i], entities_dof_end[i]): - self.dofs_info[i_d].dof_start = entities_dof_start[i] + if ti.static(self._options.batch_dofs_info): + for i_b in range(self._B): + for i_d in range(entities_dof_start[i], entities_dof_end[i]): + self.dofs_info[i_d, i_b].dof_start = entities_dof_start[i] + else: + for i_d in range(entities_dof_start[i], entities_dof_end[i]): + self.dofs_info[i_d].dof_start = entities_dof_start[i] if ti.static(self._use_hibernation): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) @@ -1053,7 +1072,8 @@ def _func_compute_mass_matrix(self): i_e = self.awake_entities[i_e_, i_b] for i in range(self.entities_info[i_e].n_links): i_l = self.entities_info[i_e].link_end - 1 - i - i_p = self.links_info[i_l].parent_idx + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + i_p = self.links_info[I_l].parent_idx if i_p != -1: self.links_state[i_p, i_b].crb_inertial = ( @@ -1075,7 +1095,8 @@ def _func_compute_mass_matrix(self): for i_b in range(self._B): for i_l_ in range(self.n_awake_links[i_b]): i_l = self.awake_links[i_l_, i_b] - l_info = self.links_info[i_l] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + l_info = self.links_info[I_l] for i_d in range(l_info.dof_start, l_info.dof_end): self.dofs_state[i_d, i_b].f_ang, self.dofs_state[i_d, i_b].f_vel = gu.inertial_mul( self.links_state[i_l, i_b].crb_pos, @@ -1098,10 +1119,11 @@ def _func_compute_mass_matrix(self): ) * self.mass_parent_mask[i_d, j_d] for i_d in range(e_info.dof_start, e_info.dof_end): + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d self.mass_mat[i_d, i_d, i_b] = ( self.mass_mat[i_d, i_d, i_b] - + self.dofs_info[i_d].armature - + self.dofs_info[i_d].damping * self._substep_dt + + self.dofs_info[I_d].armature + + self.dofs_info[I_d].damping * self._substep_dt ) for j_d in range(i_d + 1, e_info.dof_end): self.mass_mat[i_d, j_d, i_b] = self.mass_mat[j_d, i_d, i_b] @@ -1111,18 +1133,19 @@ def _func_compute_mass_matrix(self): # qDeriv += d qfrc_actuator / d qvel ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_d, i_b in ti.ndrange(self.n_dofs, self._B): + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d if self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.FORCE: pass elif self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.VELOCITY: self.mass_mat[i_d, i_d, i_b] = ( - self.mass_mat[i_d, i_d, i_b] + self.dofs_info[i_d].kv * self._substep_dt + self.mass_mat[i_d, i_d, i_b] + self.dofs_info[I_d].kv * self._substep_dt ) elif self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.POSITION: self.mass_mat[i_d, i_d, i_b] = ( - self.mass_mat[i_d, i_d, i_b] + self.dofs_info[i_d].kv * self._substep_dt + self.mass_mat[i_d, i_d, i_b] + self.dofs_info[I_d].kv * self._substep_dt ) else: @@ -1139,7 +1162,8 @@ def _func_compute_mass_matrix(self): for i_e, i_b in ti.ndrange(self.n_entities, self._B): for i in range(self.entities_info[i_e].n_links): i_l = self.entities_info[i_e].link_end - 1 - i - i_p = self.links_info[i_l].parent_idx + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + i_p = self.links_info[I_l].parent_idx if i_p != -1: self.links_state[i_p, i_b].crb_inertial = ( @@ -1160,7 +1184,8 @@ def _func_compute_mass_matrix(self): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self._B): for i_l in range(self.n_links): - l_info = self.links_info[i_l] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + l_info = self.links_info[I_l] for i_d in range(l_info.dof_start, l_info.dof_end): self.dofs_state[i_d, i_b].f_ang, self.dofs_state[i_d, i_b].f_vel = gu.inertial_mul( self.links_state[i_l, i_b].crb_pos, @@ -1181,10 +1206,11 @@ def _func_compute_mass_matrix(self): ) * self.mass_parent_mask[i_d, j_d] for i_d in range(e_info.dof_start, e_info.dof_end): + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d self.mass_mat[i_d, i_d, i_b] = ( self.mass_mat[i_d, i_d, i_b] - + self.dofs_info[i_d].armature - + self.dofs_info[i_d].damping * self._substep_dt + + self.dofs_info[I_d].armature + + self.dofs_info[I_d].damping * self._substep_dt ) for j_d in range(i_d + 1, e_info.dof_end): self.mass_mat[i_d, j_d, i_b] = self.mass_mat[j_d, i_d, i_b] @@ -1194,18 +1220,19 @@ def _func_compute_mass_matrix(self): # qDeriv += d qfrc_actuator / d qvel ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_d, i_b in ti.ndrange(self.n_dofs, self._B): + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d if self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.FORCE: pass elif self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.VELOCITY: self.mass_mat[i_d, i_d, i_b] = ( - self.mass_mat[i_d, i_d, i_b] + self.dofs_info[i_d].kv * self._substep_dt + self.mass_mat[i_d, i_d, i_b] + self.dofs_info[I_d].kv * self._substep_dt ) elif self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.POSITION: self.mass_mat[i_d, i_d, i_b] = ( - self.mass_mat[i_d, i_d, i_b] + self.dofs_info[i_d].kv * self._substep_dt + self.mass_mat[i_d, i_d, i_b] + self.dofs_info[I_d].kv * self._substep_dt ) @ti.func @@ -1353,15 +1380,16 @@ def _func_implicit_damping(self): # TODO: hibernate ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_d, i_b in ti.ndrange(self.n_dofs, self._B): + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d if self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.FORCE: pass elif self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.VELOCITY: - self.mass_mat[i_d, i_d, i_b] = self.mass_mat[i_d, i_d, i_b] + self.dofs_info[i_d].kv * self._substep_dt + self.mass_mat[i_d, i_d, i_b] = self.mass_mat[i_d, i_d, i_b] + self.dofs_info[I_d].kv * self._substep_dt elif self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.POSITION: - self.mass_mat[i_d, i_d, i_b] = self.mass_mat[i_d, i_d, i_b] + self.dofs_info[i_d].kv * self._substep_dt + self.mass_mat[i_d, i_d, i_b] = self.mass_mat[i_d, i_d, i_b] + self.dofs_info[I_d].kv * self._substep_dt self.dofs_state[i_d, i_b].force += self.dofs_state[i_d, i_b].qf_constraint @@ -1491,9 +1519,10 @@ def _func_COM_links(self): for i_b in range(self._B): for i_l_ in range(self.n_awake_links[i_b]): i_l = self.awake_links[i_l_, i_b] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l l = self.links_state[i_l, i_b] - l_info = self.links_info[i_l] + l_info = self.links_info[I_l] mass = l_info.inertial_mass + l.mass_shift self.links_state[i_l, i_b].i_pos, self.links_state[i_l, i_b].i_quat = ( gu.ti_transform_pos_quat_by_trans_quat( @@ -1501,7 +1530,7 @@ def _func_COM_links(self): ) ) - i_r = self.links_info[i_l].root_idx + i_r = self.links_info[I_l].root_idx ti.atomic_add(self.links_state[i_r, i_b].mass_sum, mass) COM = mass * self.links_state[i_l, i_b].i_pos @@ -1511,8 +1540,9 @@ def _func_COM_links(self): for i_b in range(self._B): for i_l_ in range(self.n_awake_links[i_b]): i_l = self.awake_links[i_l_, i_b] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - i_r = self.links_info[i_l].root_idx + i_r = self.links_info[I_l].root_idx if i_l == i_r: self.links_state[i_l, i_b].root_COM = ( self.links_state[i_l, i_b].root_COM / self.links_state[i_l, i_b].mass_sum @@ -1522,19 +1552,21 @@ def _func_COM_links(self): for i_b in range(self._B): for i_l_ in range(self.n_awake_links[i_b]): i_l = self.awake_links[i_l_, i_b] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - i_r = self.links_info[i_l].root_idx + i_r = self.links_info[I_l].root_idx self.links_state[i_l, i_b].root_COM = self.links_state[i_r, i_b].root_COM ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self._B): for i_l_ in range(self.n_awake_links[i_b]): i_l = self.awake_links[i_l_, i_b] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l l = self.links_state[i_l, i_b] - l_info = self.links_info[i_l] + l_info = self.links_info[I_l] - i_r = self.links_info[i_l].root_idx + i_r = self.links_info[I_l].root_idx self.links_state[i_l, i_b].COM = self.links_state[i_r, i_b].root_COM self.links_state[i_l, i_b].i_pos = self.links_state[i_l, i_b].i_pos - self.links_state[i_l, i_b].COM @@ -1553,8 +1585,9 @@ def _func_COM_links(self): for i_b in range(self._B): for i_l_ in range(self.n_awake_links[i_b]): i_l = self.awake_links[i_l_, i_b] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - l_info = self.links_info[i_l] + l_info = self.links_info[I_l] i_p = l_info.parent_idx p_pos = ti.Vector.zero(gs.ti_float, 3) @@ -1586,14 +1619,16 @@ def _func_COM_links(self): for i_b in range(self._B): for i_l_ in range(self.n_awake_links[i_b]): i_l = self.awake_links[i_l_, i_b] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - l_info = self.links_info[i_l] + l_info = self.links_info[I_l] if l_info.joint_type == gs.JOINT_TYPE.FREE: for i_d in range(l_info.dof_start, l_info.dof_end): - self.dofs_state[i_d, i_b].cdof_vel = self.dofs_info[i_d].motion_vel + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d + self.dofs_state[i_d, i_b].cdof_vel = self.dofs_info[I_d].motion_vel self.dofs_state[i_d, i_b].cdof_ang = gu.ti_transform_by_quat( - self.dofs_info[i_d].motion_ang, self.links_state[i_l, i_b].j_quat + self.dofs_info[I_d].motion_ang, self.links_state[i_l, i_b].j_quat ) offset_pos = self.links_state[i_l, i_b].COM - self.links_state[i_l, i_b].j_pos @@ -1617,8 +1652,9 @@ def _func_COM_links(self): pass else: for i_d in range(l_info.dof_start, l_info.dof_end): - motion_vel = self.dofs_info[i_d].motion_vel - motion_ang = self.dofs_info[i_d].motion_ang + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d + motion_vel = self.dofs_info[I_d].motion_vel + motion_ang = self.dofs_info[I_d].motion_ang self.dofs_state[i_d, i_b].cdof_ang = gu.ti_transform_by_quat( motion_ang, self.links_state[i_l, i_b].j_quat @@ -1654,9 +1690,10 @@ def _func_COM_links(self): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self._B): for i_l in range(self.n_links): + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l l = self.links_state[i_l, i_b] - l_info = self.links_info[i_l] + l_info = self.links_info[I_l] mass = l_info.inertial_mass + l.mass_shift self.links_state[i_l, i_b].i_pos, self.links_state[i_l, i_b].i_quat = ( gu.ti_transform_pos_quat_by_trans_quat( @@ -1664,7 +1701,7 @@ def _func_COM_links(self): ) ) - i_r = self.links_info[i_l].root_idx + i_r = self.links_info[I_l].root_idx ti.atomic_add(self.links_state[i_r, i_b].mass_sum, mass) COM = mass * self.links_state[i_l, i_b].i_pos @@ -1673,8 +1710,9 @@ def _func_COM_links(self): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self._B): for i_l in range(self.n_links): + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - i_r = self.links_info[i_l].root_idx + i_r = self.links_info[I_l].root_idx if i_l == i_r: self.links_state[i_l, i_b].root_COM = ( self.links_state[i_l, i_b].root_COM / self.links_state[i_l, i_b].mass_sum @@ -1683,18 +1721,20 @@ def _func_COM_links(self): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self._B): for i_l in range(self.n_links): + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - i_r = self.links_info[i_l].root_idx + i_r = self.links_info[I_l].root_idx self.links_state[i_l, i_b].root_COM = self.links_state[i_r, i_b].root_COM ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self._B): for i_l in range(self.n_links): + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l l = self.links_state[i_l, i_b] - l_info = self.links_info[i_l] + l_info = self.links_info[I_l] - i_r = self.links_info[i_l].root_idx + i_r = self.links_info[I_l].root_idx self.links_state[i_l, i_b].COM = self.links_state[i_r, i_b].root_COM self.links_state[i_l, i_b].i_pos = self.links_state[i_l, i_b].i_pos - self.links_state[i_l, i_b].COM @@ -1712,8 +1752,9 @@ def _func_COM_links(self): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self._B): for i_l in range(self.n_links): + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - l_info = self.links_info[i_l] + l_info = self.links_info[I_l] i_p = l_info.parent_idx p_pos = ti.Vector.zero(gs.ti_float, 3) @@ -1744,14 +1785,16 @@ def _func_COM_links(self): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self._B): for i_l in range(self.n_links): + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - l_info = self.links_info[i_l] + l_info = self.links_info[I_l] if l_info.joint_type == gs.JOINT_TYPE.FREE: for i_d in range(l_info.dof_start, l_info.dof_end): - self.dofs_state[i_d, i_b].cdof_vel = self.dofs_info[i_d].motion_vel + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d + self.dofs_state[i_d, i_b].cdof_vel = self.dofs_info[I_d].motion_vel self.dofs_state[i_d, i_b].cdof_ang = gu.ti_transform_by_quat( - self.dofs_info[i_d].motion_ang, self.links_state[i_l, i_b].j_quat + self.dofs_info[I_d].motion_ang, self.links_state[i_l, i_b].j_quat ) offset_pos = self.links_state[i_l, i_b].COM - self.links_state[i_l, i_b].j_pos @@ -1776,8 +1819,9 @@ def _func_COM_links(self): else: for i_d in range(l_info.dof_start, l_info.dof_end): - motion_vel = self.dofs_info[i_d].motion_vel - motion_ang = self.dofs_info[i_d].motion_ang + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d + motion_vel = self.dofs_info[I_d].motion_vel + motion_ang = self.dofs_info[I_d].motion_ang self.dofs_state[i_d, i_b].cdof_ang = gu.ti_transform_by_quat( motion_ang, self.links_state[i_l, i_b].j_quat @@ -1836,7 +1880,8 @@ def _func_COM_cd(self): e_info = self.entities_info[i_e] for i_l in range(e_info.link_start, e_info.link_end): - l_info = self.links_info[i_l] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + l_info = self.links_info[I_l] i_p = l_info.parent_idx cd_vel = ti.Vector.zero(gs.ti_float, 3) @@ -1861,8 +1906,9 @@ def _func_COM_cdofd(self): for i_b in range(self._B): for i_l_ in range(self.n_awake_links[i_b]): i_l = self.awake_links[i_l_, i_b] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - l_info = self.links_info[i_l] + l_info = self.links_info[I_l] if l_info.joint_type == gs.JOINT_TYPE.FREE: cd_ang = ti.Vector.zero(gs.ti_float, 3) @@ -1904,8 +1950,9 @@ def _func_COM_cdofd(self): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self._B): for i_l in range(self.n_links): + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - l_info = self.links_info[i_l] + l_info = self.links_info[I_l] if l_info.joint_type == gs.JOINT_TYPE.FREE: cd_ang = ti.Vector.zero(gs.ti_float, 3) @@ -1967,7 +2014,8 @@ def _func_forward_kinematics(self): def _func_forward_kinematics_entity(self, i_e, i_b): # calculate_j for i_l in range(self.entities_info[i_e].link_start, self.entities_info[i_e].link_end): - l_info = self.links_info[i_l] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + l_info = self.links_info[I_l] if l_info.joint_type == gs.JOINT_TYPE.FREE: @@ -1991,7 +2039,8 @@ def _func_forward_kinematics_entity(self, i_e, i_b): else: self.dofs_state[l_info.dof_start, i_b].pos = self.qpos[l_info.q_start, i_b] - dof_info = self.dofs_info[l_info.dof_start] + I_dof_start = [l_info.dof_start, i_b] if ti.static(self._options.batch_dofs_info) else l_info.dof_start + dof_info = self.dofs_info[I_dof_start] self.links_state[i_l, i_b].j_pos = dof_info.motion_vel * self.qpos[l_info.q_start, i_b] self.links_state[i_l, i_b].j_quat = gu.ti_rotvec_to_quat( dof_info.motion_ang * self.qpos[l_info.q_start, i_b] @@ -2000,7 +2049,8 @@ def _func_forward_kinematics_entity(self, i_e, i_b): self.links_state[i_l, i_b].j_vel = dof_info.motion_vel * self.dofs_state[l_info.dof_start, i_b].vel for i_d in range(l_info.dof_start + 1, l_info.dof_end): - dof_info = self.dofs_info[i_d] + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d + dof_info = self.dofs_info[I_d] qi = l_info.q_start + i_d - l_info.dof_start self.dofs_state[i_d, i_b].pos = self.qpos[qi, i_b] ji_pos = dof_info.motion_vel * self.qpos[qi, i_b] @@ -2041,7 +2091,8 @@ def _func_forward_kinematics_entity(self, i_e, i_b): # joint_to_world for i_l in range(self.entities_info[i_e].link_start, self.entities_info[i_e].link_end): l_state = self.links_state[i_l, i_b] - l_info = self.links_info[i_l] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + l_info = self.links_info[I_l] i_p = l_info.parent_idx if i_p == -1: # root link @@ -2318,15 +2369,17 @@ def _func_apply_external_torque_link_frame(self, torque, link_idx, batch_idx): @ti.func def _func_apply_external_force_link_inertial_frame(self, pos, force, link_idx, batch_idx): + link_I = [link_idx, batch_idx] if ti.static(self._options.batch_links_info) else link_idx pos = gu.ti_transform_by_trans_quat( - pos, self.links_info[link_idx].inertial_pos, self.links_info[link_idx].inertial_quat + pos, self.links_info[link_I].inertial_pos, self.links_info[link_I].inertial_quat ) - force = gu.ti_transform_by_quat(force, self.links_info[link_idx].inertial_quat) + force = gu.ti_transform_by_quat(force, self.links_info[link_I].inertial_quat) self._func_apply_external_force_link_frame(pos, force, link_idx, batch_idx) @ti.func def _func_apply_external_torque_link_inertial_frame(self, torque, link_idx, batch_idx): - torque = gu.ti_transform_by_quat(torque, self.links_info[link_idx].inertial_quat) + link_I = [link_idx, batch_idx] if ti.static(self._options.batch_links_info) else link_idx + torque = gu.ti_transform_by_quat(torque, self.links_info[link_I].inertial_quat) self._func_apply_external_torque_link_frame(torque, link_idx, batch_idx) @ti.func @@ -2352,27 +2405,29 @@ def _func_torque_and_passive_force(self): wakeup = False for i_l in range(self.entities_info[i_e].link_start, self.entities_info[i_e].link_end): force = gs.ti_float(0.0) - l_info = self.links_info[i_l] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + l_info = self.links_info[I_l] for i_d in range(l_info.dof_start, l_info.dof_end): + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d if self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.FORCE: force = self.dofs_state[i_d, i_b].ctrl_force elif self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.VELOCITY: - force = self.dofs_info[i_d].kv * ( + force = self.dofs_info[I_d].kv * ( self.dofs_state[i_d, i_b].ctrl_vel - self.dofs_state[i_d, i_b].vel ) elif self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.POSITION and not ( l_info.joint_type == gs.JOINT_TYPE.FREE and i_d >= l_info.dof_start + 3 ): force = ( - self.dofs_info[i_d].kp + self.dofs_info[I_d].kp * (self.dofs_state[i_d, i_b].ctrl_pos - self.dofs_state[i_d, i_b].pos) - - self.dofs_info[i_d].kv * self.dofs_state[i_d, i_b].vel + - self.dofs_info[I_d].kv * self.dofs_state[i_d, i_b].vel ) self.dofs_state[i_d, i_b].qf_applied = ti.math.clamp( force, - self.dofs_info[i_d].force_range[0], - self.dofs_info[i_d].force_range[1], + self.dofs_info[I_d].force_range[0], + self.dofs_info[I_d].force_range[1], ) if ti.abs(force) > gs.EPS: @@ -2410,13 +2465,14 @@ def _func_torque_and_passive_force(self): rotvec = gu.ti_quat_to_rotvec(q_diff) for i_d in range(l_info.dof_start + 3, l_info.dof_end): + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d force = ( - self.dofs_info[i_d].kp * rotvec[i_d - l_info.dof_start - 3] - - self.dofs_info[i_d].kv * self.dofs_state[i_d, i_b].vel + self.dofs_info[I_d].kp * rotvec[i_d - l_info.dof_start - 3] + - self.dofs_info[I_d].kv * self.dofs_state[i_d, i_b].vel ) self.dofs_state[i_d, i_b].qf_applied = ti.math.clamp( - force, self.dofs_info[i_d].force_range[0], self.dofs_info[i_d].force_range[1] + force, self.dofs_info[I_d].force_range[0], self.dofs_info[I_d].force_range[1] ) if ti.abs(force) > gs.EPS: wakeup = True @@ -2430,7 +2486,8 @@ def _func_torque_and_passive_force(self): for i_b in range(self._B): for i_l_ in range(self.n_awake_links[i_b]): i_l = self.awake_links[i_l_, i_b] - l_info = self.links_info[i_l] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + l_info = self.links_info[I_l] if ( l_info.joint_type == gs.JOINT_TYPE.REVOLUTE or l_info.joint_type == gs.JOINT_TYPE.PRISMATIC @@ -2442,20 +2499,23 @@ def _func_torque_and_passive_force(self): q_end = l_info.q_end for j_d in range(q_end - q_start): + I_d = [dof_start + j_d, i_b] if ti.static(self._options.batch_dofs_info) else dof_start + j_d self.dofs_state[dof_start + j_d, i_b].qf_passive = ( - -self.qpos[q_start + j_d, i_b] * self.dofs_info[dof_start + j_d].stiffness + -self.qpos[q_start + j_d, i_b] * self.dofs_info[I_d].stiffness ) ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_b in range(self._B): for i_d_ in range(self.n_awake_dofs[i_b]): i_d = self.awake_dofs[i_d_, i_b] + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d - self.dofs_state[i_d, i_b].qf_passive += -self.dofs_info[i_d].damping * self.dofs_state[i_d, i_b].vel + self.dofs_state[i_d, i_b].qf_passive += -self.dofs_info[I_d].damping * self.dofs_state[i_d, i_b].vel else: ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_l, i_b in ti.ndrange(self.n_links, self._B): - l_info = self.links_info[i_l] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + l_info = self.links_info[I_l] if ( l_info.joint_type == gs.JOINT_TYPE.REVOLUTE or l_info.joint_type == gs.JOINT_TYPE.PRISMATIC @@ -2467,13 +2527,15 @@ def _func_torque_and_passive_force(self): q_end = l_info.q_end for j_d in range(q_end - q_start): + I_d = [dof_start + j_d, i_b] if ti.static(self._options.batch_dofs_info) else dof_start + j_d self.dofs_state[dof_start + j_d, i_b].qf_passive = ( - -self.qpos[q_start + j_d, i_b] * self.dofs_info[dof_start + j_d].stiffness + -self.qpos[q_start + j_d, i_b] * self.dofs_info[I_d].stiffness ) ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_d, i_b in ti.ndrange(self.n_dofs, self._B): - self.dofs_state[i_d, i_b].qf_passive += -self.dofs_info[i_d].damping * self.dofs_state[i_d, i_b].vel + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d + self.dofs_state[i_d, i_b].qf_passive += -self.dofs_info[I_d].damping * self.dofs_state[i_d, i_b].vel @ti.func def _func_system_update_acc(self): @@ -2484,8 +2546,9 @@ def _func_system_update_acc(self): i_e = self.awake_entities[i_e_, i_b] e_info = self.entities_info[i_e] for i_l in range(e_info.link_start, e_info.link_end): + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - i_p = self.links_info[i_l].parent_idx + i_p = self.links_info[I_l].parent_idx if i_p == -1: self.links_state[i_l, i_b].cdd_vel = -self._gravity[None] * ( 1 - e_info.gravity_compensation @@ -2499,7 +2562,7 @@ def _func_system_update_acc(self): map_sum_vel = ti.Vector.zero(gs.ti_float, 3) map_sum_ang = ti.Vector.zero(gs.ti_float, 3) - for i_d in range(self.links_info[i_l].dof_start, self.links_info[i_l].dof_end): + for i_d in range(self.links_info[I_l].dof_start, self.links_info[I_l].dof_end): map_sum_vel = ( map_sum_vel + self.dofs_state[i_d, i_b].cdofd_vel * self.dofs_state[i_d, i_b].vel ) @@ -2514,8 +2577,9 @@ def _func_system_update_acc(self): for i_e, i_b in ti.ndrange(self.n_entities, self._B): e_info = self.entities_info[i_e] for i_l in range(e_info.link_start, e_info.link_end): + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - i_p = self.links_info[i_l].parent_idx + i_p = self.links_info[I_l].parent_idx if i_p == -1: self.links_state[i_l, i_b].cdd_vel = -self._gravity[None] * (1 - e_info.gravity_compensation) self.links_state[i_l, i_b].cdd_ang = ti.Vector.zero(gs.ti_float, 3) @@ -2527,7 +2591,7 @@ def _func_system_update_acc(self): map_sum_vel = ti.Vector.zero(gs.ti_float, 3) map_sum_ang = ti.Vector.zero(gs.ti_float, 3) - for i_d in range(self.links_info[i_l].dof_start, self.links_info[i_l].dof_end): + for i_d in range(self.links_info[I_l].dof_start, self.links_info[I_l].dof_end): map_sum_vel = map_sum_vel + self.dofs_state[i_d, i_b].cdofd_vel * self.dofs_state[i_d, i_b].vel map_sum_ang = map_sum_ang + self.dofs_state[i_d, i_b].cdofd_ang * self.dofs_state[i_d, i_b].vel @@ -2596,7 +2660,8 @@ def _func_inverse_link_force(self): e_info = self.entities_info[i_e] for i in range(e_info.n_links): i_l = e_info.link_end - 1 - i - i_p = self.links_info[i_l].parent_idx + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + i_p = self.links_info[I_l].parent_idx if i_p != -1: self.links_state[i_p, i_b].cfrc_flat_vel = ( self.links_state[i_p, i_b].cfrc_flat_vel + self.links_state[i_l, i_b].cfrc_flat_vel @@ -2611,7 +2676,8 @@ def _func_inverse_link_force(self): e_info = self.entities_info[i_e] for i in range(e_info.n_links): i_l = e_info.link_end - 1 - i - i_p = self.links_info[i_l].parent_idx + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + i_p = self.links_info[I_l].parent_idx if i_p != -1: self.links_state[i_p, i_b].cfrc_flat_vel = ( self.links_state[i_p, i_b].cfrc_flat_vel + self.links_state[i_l, i_b].cfrc_flat_vel @@ -2628,15 +2694,16 @@ def _func_actuation(self): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_l, i_b in ti.ndrange(self.n_links, self._B): joint_type = self.links_info[i_l].joint_type - q_start = self.links_info[i_l].q_start + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + q_start = self.links_info[I_l].q_start if joint_type == gs.JOINT_TYPE.REVOLUTE or joint_type == gs.JOINT_TYPE.PRISMATIC: gear = -1 # TODO - i_d = self.links_info[i_l].dof_start + i_d = self.links_info[I_l].dof_start self.dofs_state[i_d, i_b].act_length = gear * self.qpos[q_start, i_b] self.dofs_state[i_d, i_b].qf_actuator = self.dofs_state[i_d, i_b].act_length else: - for i_d in range(self.links_info[i_l].dof_start, self.links_info[i_l].dof_end): + for i_d in range(self.links_info[I_l].dof_start, self.links_info[I_l].dof_end): self.dofs_state[i_d, i_b].act_length = 0.0 self.dofs_state[i_d, i_b].qf_actuator = self.dofs_state[i_d, i_b].act_length @@ -2647,8 +2714,9 @@ def _func_bias_force(self): for i_b in range(self._B): for i_l_ in range(self.n_awake_links[i_b]): i_l = self.awake_links[i_l_, i_b] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - l_info = self.links_info[i_l] + l_info = self.links_info[I_l] dof_start = l_info.dof_start dof_end = l_info.dof_end @@ -2667,8 +2735,9 @@ def _func_bias_force(self): else: ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_l, i_b in ti.ndrange(self.n_links, self._B): + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l - l_info = self.links_info[i_l] + l_info = self.links_info[I_l] dof_start = l_info.dof_start dof_end = l_info.dof_end @@ -2727,10 +2796,11 @@ def _func_integrate(self): for i_b in range(self._B): for i_l_ in range(self.n_awake_links[i_b]): i_l = self.awake_links[i_l_, i_b] - joint_type = self.links_info[i_l].joint_type - dof_start = self.links_info[i_l].dof_start - q_start = self.links_info[i_l].q_start - q_end = self.links_info[i_l].q_end + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + joint_type = self.links_info[I_l].joint_type + dof_start = self.links_info[I_l].dof_start + q_start = self.links_info[I_l].q_start + q_end = self.links_info[I_l].q_end if joint_type == gs.JOINT_TYPE.FREE: rot = ti.Vector( @@ -2795,10 +2865,11 @@ def _func_integrate(self): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.ALL) for i_l, i_b in ti.ndrange(self.n_links, self._B): - joint_type = self.links_info[i_l].joint_type - dof_start = self.links_info[i_l].dof_start - q_start = self.links_info[i_l].q_start - q_end = self.links_info[i_l].q_end + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + joint_type = self.links_info[I_l].joint_type + dof_start = self.links_info[I_l].dof_start + q_start = self.links_info[I_l].q_start + q_end = self.links_info[I_l].q_end if joint_type == gs.JOINT_TYPE.FREE: rot = ti.Vector( @@ -2844,7 +2915,8 @@ def _func_integrate(self): def _func_integrate_dq_entity(self, dq, i_e, i_b, respect_joint_limit): e_info = self.entities_info[i_e] for i_l in range(e_info.link_start, e_info.link_end): - l_info = self.links_info[i_l] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + l_info = self.links_info[I_l] q_start = l_info.q_start dof_start = l_info.dof_start dq_start = l_info.dof_start - e_info.dof_start @@ -2883,10 +2955,11 @@ def _func_integrate_dq_entity(self, dq, i_e, i_b, respect_joint_limit): self.qpos[q_start + i_d_, i_b] = self.qpos[q_start + i_d_, i_b] + dq[dq_start + i_d_, i_b] if respect_joint_limit: + I_d = [dof_start + i_d_, i_b] if ti.static(self._options.batch_dofs_info) else dof_start + i_d_ self.qpos[q_start + i_d_, i_b] = ti.math.clamp( self.qpos[q_start + i_d_, i_b], - self.dofs_info[dof_start + i_d_].limit[0], - self.dofs_info[dof_start + i_d_].limit[1], + self.dofs_info[I_d].limit[0], + self.dofs_info[I_d].limit[1], ) def substep_pre_coupling(self, f): @@ -3206,12 +3279,13 @@ def _kernel_set_links_pos( ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]): i_l = links_idx[i_l_] - if self.links_info[i_l].is_fixed: # change links_state directly as the link's pose is not contained in qpos + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + if self.links_info[I_l].is_fixed: # change links_state directly as the link's pose is not contained in qpos for i in ti.static(range(3)): self.links_state[i_l, envs_idx[i_b_]].pos[i] = pos[i_b_, i_l_, i] else: # free base link's pose is reflected in qpos, and links_state will be computed automatically - q_start = self.links_info[i_l].q_start + q_start = self.links_info[I_l].q_start for i in ti.static(range(3)): self.qpos[q_start + i, envs_idx[i_b_]] = pos[i_b_, i_l_, i] @@ -3239,12 +3313,13 @@ def _kernel_set_links_quat( ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]): i_l = links_idx[i_l_] - if self.links_info[i_l].is_fixed: # change links_state directly as the link's pose is not contained in qpos + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + if self.links_info[I_l].is_fixed: # change links_state directly as the link's pose is not contained in qpos for i in ti.static(range(4)): self.links_state[i_l, envs_idx[i_b_]].quat[i] = quat[i_b_, i_l_, i] else: # free base link's pose is reflected in qpos, and links_state will be computed automatically - q_start = self.links_info[i_l].q_start + q_start = self.links_info[I_l].q_start for i in ti.static(range(4)): self.qpos[q_start + i + 3, envs_idx[i_b_]] = quat[i_b_, i_l_, i] @@ -3337,10 +3412,11 @@ def _kernel_set_global_sol_params(self, sol_params: ti.types.ndarray()): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i, b in ti.ndrange(self.n_dofs, self._B): + I = [i, b] if ti.static(self._options.batch_dofs_info) else i for j in ti.static(range(7)): - self.dofs_info[i].sol_params[j] = sol_params[j] + self.dofs_info[I].sol_params[j] = sol_params[j] - self.dofs_info[i].sol_params[0] = self._substep_dt * 2 + self.dofs_info[I].sol_params[0] = self._substep_dt * 2 def set_dofs_kp(self, kp, dofs_idx): kp, dofs_idx = self._validate_1D_io_variables(kp, dofs_idx, batched=False) @@ -3351,6 +3427,7 @@ def _kernel_set_dofs_kp( self, kp: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), + # TODO: batch ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i_d_ in range(dofs_idx.shape[0]): @@ -3365,6 +3442,7 @@ def _kernel_set_dofs_kv( self, kv: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), + # TODO: batch ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i_d_ in range(dofs_idx.shape[0]): @@ -3385,6 +3463,7 @@ def _kernel_set_dofs_force_range( lower: ti.types.ndarray(), upper: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), + # TODO: batch ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i_d_ in range(dofs_idx.shape[0]): @@ -3437,7 +3516,8 @@ def _kernel_set_dofs_position( for i_e, i_b_ in ti.ndrange(self.n_entities, envs_idx.shape[0]): i_b = envs_idx[i_b_] for i_l in range(self.entities_info[i_e].link_start, self.entities_info[i_e].link_end): - l_info = self.links_info[i_l] + I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + l_info = self.links_info[I_l] if l_info.joint_type == gs.JOINT_TYPE.FREE: xyz = ti.Vector( @@ -3767,20 +3847,21 @@ def _kernel_get_dofs_control_force( for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): i_d = dofs_idx[i_d_] i_b = envs_idx[i_b_] + I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d force = gs.ti_float(0.0) if self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.FORCE: force = self.dofs_state[i_d, i_b].ctrl_force elif self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.VELOCITY: - force = self.dofs_info[i_d].kv * (self.dofs_state[i_d, i_b].ctrl_vel - self.dofs_state[i_d, i_b].vel) + force = self.dofs_info[I_d].kv * (self.dofs_state[i_d, i_b].ctrl_vel - self.dofs_state[i_d, i_b].vel) elif self.dofs_state[i_d, i_b].ctrl_mode == gs.CTRL_MODE.POSITION: force = ( - self.dofs_info[i_d].kp * (self.dofs_state[i_d, i_b].ctrl_pos - self.dofs_state[i_d, i_b].pos) - - self.dofs_info[i_d].kv * self.dofs_state[i_d, i_b].vel + self.dofs_info[I_d].kp * (self.dofs_state[i_d, i_b].ctrl_pos - self.dofs_state[i_d, i_b].pos) + - self.dofs_info[I_d].kv * self.dofs_state[i_d, i_b].vel ) tensor[i_b_, i_d_] = ti.math.clamp( force, - self.dofs_info[i_d].force_range[0], - self.dofs_info[i_d].force_range[1], + self.dofs_info[I_d].force_range[0], + self.dofs_info[I_d].force_range[1], ) @ti.kernel @@ -3859,6 +3940,7 @@ def _kernel_get_dofs_kp( self, tensor: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), + # TODO: batch ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i_d_ in range(dofs_idx.shape[0]): @@ -3869,6 +3951,7 @@ def _kernel_get_dofs_kv( self, tensor: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), + # TODO: batch ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i_d_ in range(dofs_idx.shape[0]): @@ -3880,6 +3963,7 @@ def _kernel_get_dofs_force_range( lower: ti.types.ndarray(), upper: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), + # TODO: batch ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i_d_ in range(dofs_idx.shape[0]): @@ -3892,6 +3976,7 @@ def _kernel_get_dofs_limit( lower: ti.types.ndarray(), upper: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), + # TODO: batch ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i_d_ in range(dofs_idx.shape[0]): diff --git a/genesis/options/solvers.py b/genesis/options/solvers.py index 54e80b18..1b3b6c51 100644 --- a/genesis/options/solvers.py +++ b/genesis/options/solvers.py @@ -185,6 +185,10 @@ class RigidOptions(Options): integrator: gs.integrator = gs.integrator.approximate_implicitfast IK_max_targets: int = 6 + # batching info + batch_links_info: Optional[bool] = False + batch_dofs_info: Optional[bool] = False + # constraint solver constraint_solver: gs.constraint_solver = gs.constraint_solver.CG iterations: int = 100 From af9f052e099393bac05ed6a1f12b2df8d33c10ff Mon Sep 17 00:00:00 2001 From: Johnson Wang Date: Sun, 22 Dec 2024 07:27:17 -0500 Subject: [PATCH 04/35] fix existing api for batching info; fix minor in collider --- .../entities/rigid_entity/rigid_entity.py | 54 +++++--- .../engine/solvers/rigid/collider_decomp.py | 4 + .../solvers/rigid/rigid_solver_decomp.py | 126 ++++++++++++------ 3 files changed, 125 insertions(+), 59 deletions(-) diff --git a/genesis/engine/entities/rigid_entity/rigid_entity.py b/genesis/engine/entities/rigid_entity/rigid_entity.py index ce53344a..82f6df00 100644 --- a/genesis/engine/entities/rigid_entity/rigid_entity.py +++ b/genesis/engine/entities/rigid_entity/rigid_entity.py @@ -1722,7 +1722,7 @@ def set_qpos(self, qpos, qs_idx_local=None, zero_velocity=True, envs_idx=None): self.zero_all_dofs_velocity(envs_idx) @gs.assert_built - def set_dofs_kp(self, kp, dofs_idx_local=None): + def set_dofs_kp(self, kp, dofs_idx_local=None, envs_idx=None): """ Set the entity's dofs' positional gains for the PD controller. @@ -1732,12 +1732,14 @@ def set_dofs_kp(self, kp, dofs_idx_local=None): The positional gains to set. dofs_idx_local : None | array_like, optional The indices of the dofs to set. If None, all dofs will be set. Note that here this uses the local `q_idx`, not the scene-level one. Defaults to None. + envs_idx : None | array_like, optional + The indices of the environments. If None, all environments will be considered. Defaults to None. """ - self._solver.set_dofs_kp(kp, self._get_dofs_idx(dofs_idx_local)) + self._solver.set_dofs_kp(kp, self._get_dofs_idx(dofs_idx_local), envs_idx) @gs.assert_built - def set_dofs_kv(self, kv, dofs_idx_local=None): + def set_dofs_kv(self, kv, dofs_idx_local=None, envs_idx=None): """ Set the entity's dofs' velocity gains for the PD controller. @@ -1747,11 +1749,13 @@ def set_dofs_kv(self, kv, dofs_idx_local=None): The velocity gains to set. dofs_idx_local : None | array_like, optional The indices of the dofs to set. If None, all dofs will be set. Note that here this uses the local `q_idx`, not the scene-level one. Defaults to None. + envs_idx : None | array_like, optional + The indices of the environments. If None, all environments will be considered. Defaults to None. """ - self._solver.set_dofs_kv(kv, self._get_dofs_idx(dofs_idx_local)) + self._solver.set_dofs_kv(kv, self._get_dofs_idx(dofs_idx_local), envs_idx) @gs.assert_built - def set_dofs_force_range(self, lower, upper, dofs_idx_local=None): + def set_dofs_force_range(self, lower, upper, dofs_idx_local=None, envs_idx=None): """ Set the entity's dofs' force range. @@ -1763,9 +1767,11 @@ def set_dofs_force_range(self, lower, upper, dofs_idx_local=None): The upper bounds of the force range. dofs_idx_local : None | array_like, optional The indices of the dofs to set. If None, all dofs will be set. Note that here this uses the local `q_idx`, not the scene-level one. Defaults to None. + envs_idx : None | array_like, optional + The indices of the environments. If None, all environments will be considered. Defaults to None. """ - self._solver.set_dofs_force_range(lower, upper, self._get_dofs_idx(dofs_idx_local)) + self._solver.set_dofs_force_range(lower, upper, self._get_dofs_idx(dofs_idx_local), envs_idx) @gs.assert_built def set_dofs_velocity(self, velocity, dofs_idx_local=None, envs_idx=None): @@ -1954,7 +1960,7 @@ def get_dofs_position(self, dofs_idx_local=None, envs_idx=None): return self._solver.get_dofs_position(self._get_dofs_idx(dofs_idx_local), envs_idx) @gs.assert_built - def get_dofs_kp(self, dofs_idx_local=None): + def get_dofs_kp(self, dofs_idx_local=None, envs_idx=None): """ Get the positional gain (kp) for the entity's dofs used by the PD controller. @@ -1962,16 +1968,18 @@ def get_dofs_kp(self, dofs_idx_local=None): ---------- dofs_idx_local : None | array_like, optional The indices of the dofs to get. If None, all dofs will be returned. Note that here this uses the local `q_idx`, not the scene-level one. Defaults to None. + envs_idx : None | array_like, optional + The indices of the environments. If None, all environments will be considered. Defaults to None. Returns ------- - kp : torch.Tensor, shape (n_dofs,) + kp : torch.Tensor, shape (n_dofs,) or (n_envs, n_dofs) The positional gain (kp) for the entity's dofs. """ - return self._solver.get_dofs_kp(self._get_dofs_idx(dofs_idx_local)) + return self._solver.get_dofs_kp(self._get_dofs_idx(dofs_idx_local), envs_idx) @gs.assert_built - def get_dofs_kv(self, dofs_idx_local=None): + def get_dofs_kv(self, dofs_idx_local=None, envs_idx=None): """ Get the velocity gain (kv) for the entity's dofs used by the PD controller. @@ -1979,16 +1987,18 @@ def get_dofs_kv(self, dofs_idx_local=None): ---------- dofs_idx_local : None | array_like, optional The indices of the dofs to get. If None, all dofs will be returned. Note that here this uses the local `q_idx`, not the scene-level one. Defaults to None. + envs_idx : None | array_like, optional + The indices of the environments. If None, all environments will be considered. Defaults to None. Returns ------- - kv : torch.Tensor, shape (n_dofs,) + kv : torch.Tensor, shape (n_dofs,) or (n_envs, n_dofs) The velocity gain (kv) for the entity's dofs. """ - return self._solver.get_dofs_kv(self._get_dofs_idx(dofs_idx_local)) + return self._solver.get_dofs_kv(self._get_dofs_idx(dofs_idx_local), envs_idx) @gs.assert_built - def get_dofs_force_range(self, dofs_idx_local=None): + def get_dofs_force_range(self, dofs_idx_local=None, envs_idx=None): """ Get the force range (min and max limits) for the entity's dofs. @@ -1996,18 +2006,20 @@ def get_dofs_force_range(self, dofs_idx_local=None): ---------- dofs_idx_local : None | array_like, optional The indices of the dofs to get. If None, all dofs will be returned. Note that here this uses the local `q_idx`, not the scene-level one. Defaults to None. + envs_idx : None | array_like, optional + The indices of the environments. If None, all environments will be considered. Defaults to None. Returns ------- - lower_limit : torch.Tensor, shape (n_dofs,) + lower_limit : torch.Tensor, shape (n_dofs,) or (n_envs, n_dofs) The lower limit of the force range for the entity's dofs. - upper_limit : torch.Tensor, shape (n_dofs,) + upper_limit : torch.Tensor, shape (n_dofs,) or (n_envs, n_dofs) The upper limit of the force range for the entity's dofs. """ - return self._solver.get_dofs_force_range(self._get_dofs_idx(dofs_idx_local)) + return self._solver.get_dofs_force_range(self._get_dofs_idx(dofs_idx_local), envs_idx) @gs.assert_built - def get_dofs_limit(self, dofs_idx=None): + def get_dofs_limit(self, dofs_idx=None, envs_idx=None): """ Get the positional limits (min and max) for the entity's dofs. @@ -2015,15 +2027,17 @@ def get_dofs_limit(self, dofs_idx=None): ---------- dofs_idx : None | array_like, optional The indices of the dofs to get. If None, all dofs will be returned. Note that here this uses the local `q_idx`, not the scene-level one. Defaults to None. + envs_idx : None | array_like, optional + The indices of the environments. If None, all environments will be considered. Defaults to None. Returns ------- - lower_limit : torch.Tensor, shape (n_dofs,) + lower_limit : torch.Tensor, shape (n_dofs,) or (n_envs, n_dofs) The lower limit of the positional limits for the entity's dofs. - upper_limit : torch.Tensor, shape (n_dofs,) + upper_limit : torch.Tensor, shape (n_dofs,) or (n_envs, n_dofs) The upper limit of the positional limits for the entity's dofs. """ - return self._solver.get_dofs_limit(self._get_dofs_idx(dofs_idx)) + return self._solver.get_dofs_limit(self._get_dofs_idx(dofs_idx), envs_idx) @gs.assert_built def zero_all_dofs_velocity(self, envs_idx=None): diff --git a/genesis/engine/solvers/rigid/collider_decomp.py b/genesis/engine/solvers/rigid/collider_decomp.py index 6b438750..35a7756f 100644 --- a/genesis/engine/solvers/rigid/collider_decomp.py +++ b/genesis/engine/solvers/rigid/collider_decomp.py @@ -51,6 +51,10 @@ def _init_collision_fields(self): links_root_idx = self._solver.links_info.root_idx.to_numpy() links_parent_idx = self._solver.links_info.parent_idx.to_numpy() links_is_fixed = self._solver.links_info.is_fixed.to_numpy() + if self._solver._options.batch_links_info: + links_root_idx = links_root_idx[:, 0] + links_parent_idx = links_parent_idx[:, 0] + links_is_fixed = links_is_fixed[:, 0] n_possible_pairs = 0 for i in range(self._solver.n_geoms): for j in range(i + 1, self._solver.n_geoms): diff --git a/genesis/engine/solvers/rigid/rigid_solver_decomp.py b/genesis/engine/solvers/rigid/rigid_solver_decomp.py index 985c441f..b045f694 100644 --- a/genesis/engine/solvers/rigid/rigid_solver_decomp.py +++ b/genesis/engine/solvers/rigid/rigid_solver_decomp.py @@ -3418,44 +3418,65 @@ def _kernel_set_global_sol_params(self, sol_params: ti.types.ndarray()): self.dofs_info[I].sol_params[0] = self._substep_dt * 2 - def set_dofs_kp(self, kp, dofs_idx): - kp, dofs_idx = self._validate_1D_io_variables(kp, dofs_idx, batched=False) - self._kernel_set_dofs_kp(kp, dofs_idx) + def set_dofs_kp(self, kp, dofs_idx, envs_idx=None): + if self._options.batch_dofs_info: + kp, dofs_idx, envs_idx = self._validate_1D_io_variables(kp, dofs_idx, envs_idx) + else: + kp, dofs_idx = self._validate_1D_io_variables(kp, dofs_idx, batched=False) + envs_idx = torch.empty(()) + self._kernel_set_dofs_kp(kp, dofs_idx, envs_idx) @ti.kernel def _kernel_set_dofs_kp( self, kp: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), - # TODO: batch + envs_idx: ti.types.ndarray(), ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) - for i_d_ in range(dofs_idx.shape[0]): - self.dofs_info[dofs_idx[i_d_]].kp = kp[i_d_] + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].kp = kp[i_b_, i_d_] + else: + for i_d_ in range(dofs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_]].kp = kp[i_d_] - def set_dofs_kv(self, kv, dofs_idx): - kv, dofs_idx = self._validate_1D_io_variables(kv, dofs_idx, batched=False) - self._kernel_set_dofs_kv(kv, dofs_idx) + def set_dofs_kv(self, kv, dofs_idx, envs_idx=None): + if self._options.batch_dofs_info: + kv, dofs_idx, envs_idx = self._validate_1D_io_variables(kv, dofs_idx, envs_idx) + else: + kv, dofs_idx = self._validate_1D_io_variables(kv, dofs_idx, batched=False) + envs_idx = torch.empty(()) + self._kernel_set_dofs_kv(kv, dofs_idx, envs_idx) @ti.kernel def _kernel_set_dofs_kv( self, kv: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), - # TODO: batch + envs_idx: ti.types.ndarray(), ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) - for i_d_ in range(dofs_idx.shape[0]): - self.dofs_info[dofs_idx[i_d_]].kv = kv[i_d_] + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].kv = kv[i_b_, i_d_] + else: + for i_d_ in range(dofs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_]].kv = kv[i_d_] - def set_dofs_force_range(self, lower, upper, dofs_idx): - lower, _ = self._validate_1D_io_variables(lower, dofs_idx, batched=False) - upper, dofs_idx = self._validate_1D_io_variables(upper, dofs_idx, batched=False) + def set_dofs_force_range(self, lower, upper, dofs_idx, envs_idx=None): + if self._options.batch_dofs_info: + lower, _, _ = self._validate_1D_io_variables(lower, dofs_idx, envs_idx) + upper, dofs_idx, envs_idx = self._validate_1D_io_variables(upper, dofs_idx, envs_idx) + else: + lower, _ = self._validate_1D_io_variables(lower, dofs_idx, batched=False) + upper, dofs_idx = self._validate_1D_io_variables(upper, dofs_idx, batched=False) + envs_idx = torch.empty(()) if (lower > upper).any(): gs.raise_exception("`lower` should be less than or equal to `upper`.") - self._kernel_set_dofs_force_range(lower, upper, dofs_idx) + self._kernel_set_dofs_force_range(lower, upper, dofs_idx, envs_idx) @ti.kernel def _kernel_set_dofs_force_range( @@ -3463,12 +3484,17 @@ def _kernel_set_dofs_force_range( lower: ti.types.ndarray(), upper: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), - # TODO: batch + envs_idx: ti.types.ndarray(), ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) - for i_d_ in range(dofs_idx.shape[0]): - self.dofs_info[dofs_idx[i_d_]].force_range[0] = lower[i_d_] - self.dofs_info[dofs_idx[i_d_]].force_range[1] = upper[i_d_] + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].force_range[0] = lower[i_b_, i_d_] + self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].force_range[1] = upper[i_b_, i_d_] + else: + for i_d_ in range(dofs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_]].force_range[0] = lower[i_d_] + self.dofs_info[dofs_idx[i_d_]].force_range[1] = upper[i_d_] def set_dofs_velocity(self, velocity, dofs_idx, envs_idx=None): velocity, dofs_idx, envs_idx = self._validate_1D_io_variables(velocity, dofs_idx, envs_idx) @@ -3910,26 +3936,30 @@ def get_dofs_limit(self, dofs_idx, envs_idx=None): return self._get_dofs_info(dofs_idx, "limit", envs_idx) def _get_dofs_info(self, dofs_idx, name, envs_idx=None): - tensor, dofs_idx = self._validate_1D_io_variables(None, dofs_idx, batched=False) + if self._options.batch_dofs_info: + tensor, dofs_idx, envs_idx = self._validate_1D_io_variables(None, dofs_idx, envs_idx) + else: + tensor, dofs_idx = self._validate_1D_io_variables(None, dofs_idx, batched=False) + envs_idx = torch.empty(()) if name == "kp": - self._kernel_get_dofs_kp(tensor, dofs_idx) + self._kernel_get_dofs_kp(tensor, dofs_idx, envs_idx) return tensor elif name == "kv": - self._kernel_get_dofs_kv(tensor, dofs_idx) + self._kernel_get_dofs_kv(tensor, dofs_idx, envs_idx) return tensor elif name == "force_range": lower = torch.empty_like(tensor) upper = torch.empty_like(tensor) - self._kernel_get_dofs_force_range(lower, upper, dofs_idx) + self._kernel_get_dofs_force_range(lower, upper, dofs_idx, envs_idx) return lower, upper elif name == "limit": lower = torch.empty_like(tensor) upper = torch.empty_like(tensor) - self._kernel_get_dofs_limit(lower, upper, dofs_idx) + self._kernel_get_dofs_limit(lower, upper, dofs_idx, envs_idx) return lower, upper else: @@ -3940,22 +3970,30 @@ def _kernel_get_dofs_kp( self, tensor: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), - # TODO: batch + envs_idx: ti.types.ndarray(), ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) - for i_d_ in range(dofs_idx.shape[0]): - tensor[i_d_] = self.dofs_info[dofs_idx[i_d_]].kp + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + tensor[i_b_, i_d_] = self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].kp + else: + for i_d_ in range(dofs_idx.shape[0]): + tensor[i_d_] = self.dofs_info[dofs_idx[i_d_]].kp @ti.kernel def _kernel_get_dofs_kv( self, tensor: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), - # TODO: batch + envs_idx: ti.types.ndarray(), ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) - for i_d_ in range(dofs_idx.shape[0]): - tensor[i_d_] = self.dofs_info[dofs_idx[i_d_]].kv + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + tensor[i_b_, i_d_] = self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].kv + else: + for i_d_ in range(dofs_idx.shape[0]): + tensor[i_d_] = self.dofs_info[dofs_idx[i_d_]].kv @ti.kernel def _kernel_get_dofs_force_range( @@ -3963,12 +4001,17 @@ def _kernel_get_dofs_force_range( lower: ti.types.ndarray(), upper: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), - # TODO: batch + envs_idx: ti.types.ndarray(), ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) - for i_d_ in range(dofs_idx.shape[0]): - lower[i_d_] = self.dofs_info[dofs_idx[i_d_]].force_range[0] - upper[i_d_] = self.dofs_info[dofs_idx[i_d_]].force_range[1] + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + lower[i_b_, i_d_] = self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].force_range[0] + upper[i_b_, i_d_] = self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].force_range[1] + else: + for i_d_ in range(dofs_idx.shape[0]): + lower[i_d_] = self.dofs_info[dofs_idx[i_d_]].force_range[0] + upper[i_d_] = self.dofs_info[dofs_idx[i_d_]].force_range[1] @ti.kernel def _kernel_get_dofs_limit( @@ -3976,12 +4019,17 @@ def _kernel_get_dofs_limit( lower: ti.types.ndarray(), upper: ti.types.ndarray(), dofs_idx: ti.types.ndarray(), - # TODO: batch + envs_idx: ti.types.ndarray(), ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) - for i_d_ in range(dofs_idx.shape[0]): - lower[i_d_] = self.dofs_info[dofs_idx[i_d_]].limit[0] - upper[i_d_] = self.dofs_info[dofs_idx[i_d_]].limit[1] + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + lower[i_b_, i_d_] = self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].limit[0] + upper[i_b_, i_d_] = self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].limit[1] + else: + for i_d_ in range(dofs_idx.shape[0]): + lower[i_d_] = self.dofs_info[dofs_idx[i_d_]].limit[0] + upper[i_d_] = self.dofs_info[dofs_idx[i_d_]].limit[1] @ti.kernel def _kernel_set_drone_rpm( From bded83da7e8e284e46daea2d93e5062616418e26 Mon Sep 17 00:00:00 2001 From: Johnson Wang Date: Sun, 22 Dec 2024 09:37:06 -0500 Subject: [PATCH 05/35] apis to set more attrs --- examples/rigid/set_phys_attr.py | 205 +++++++++++ .../entities/rigid_entity/rigid_entity.py | 60 ++++ .../solvers/rigid/rigid_solver_decomp.py | 340 ++++++++++++++++-- 3 files changed, 583 insertions(+), 22 deletions(-) create mode 100644 examples/rigid/set_phys_attr.py diff --git a/examples/rigid/set_phys_attr.py b/examples/rigid/set_phys_attr.py new file mode 100644 index 00000000..6979c5ec --- /dev/null +++ b/examples/rigid/set_phys_attr.py @@ -0,0 +1,205 @@ +import argparse + +import numpy as np + +import genesis as gs + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("-v", "--vis", action="store_true", default=False) + args = parser.parse_args() + + ########################## init ########################## + gs.init(backend=gs.gpu) + + ########################## create a scene ########################## + viewer_options = gs.options.ViewerOptions( + camera_pos=(0, -3.5, 2.5), + camera_lookat=(0.0, 0.0, 0.5), + camera_fov=40, + max_FPS=60, + ) + + scene = gs.Scene( + viewer_options=viewer_options, + sim_options=gs.options.SimOptions( + dt=0.01, + ), + show_viewer=args.vis, + rigid_options=gs.options.RigidOptions( + # NOTE: Batching dofs/links info to set different physical attributes across environments (in parallel) + # By default, both are False as it's faster and thus only turn this on if necessary + batch_dofs_info=True, + batch_links_info=True, + ), + ) + + ########################## entities ########################## + plane = scene.add_entity( + gs.morphs.Plane(), + ) + franka = scene.add_entity( + gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"), + ) + ########################## build ########################## + scene.build(n_envs=2) # test with 2 different environments + + jnt_names = [ + "joint1", + "joint2", + "joint3", + "joint4", + "joint5", + "joint6", + "joint7", + "finger_joint1", + "finger_joint2", + ] + dofs_idx = [franka.get_joint(name).dof_idx_local for name in jnt_names] + + lnk_names = [ + "link0", + 'link1', + 'link2', + 'link3', + 'link4', + 'link5', + 'link6', + 'link7', + 'hand', + 'left_finger', + 'right_finger', + ] + links_idx = [franka.get_link(name).idx_local for name in lnk_names] + + # Optional: set control gains + franka.set_dofs_kp( + np.array([ + [4500, 4500, 3500, 3500, 2000, 2000, 2000, 100, 100], + [100, 100, 2000, 2000, 2000, 3500, 3500, 4500, 4500], + ]), + dofs_idx, + ) + print("=== kp ===\n", franka.get_dofs_kp()) + franka.set_dofs_kv( + np.array([ + [450, 450, 350, 350, 200, 200, 200, 10, 10], + [10, 10, 200, 200, 200, 350, 350, 450, 450], + ]), + dofs_idx, + ) + print("=== kv ===\n", franka.get_dofs_kv()) + franka.set_dofs_force_range( + np.array([ + [-87, -87, -87, -87, -12, -12, -12, -100, -100], + [-120, -100, -12, -12, -12, -87, -87, -87, -87], + ]), + np.array([ + [87, 87, 87, 87, 12, 12, 12, 100, 100], + [100, 100, 12, 12, 12, 87, 87, 87, 87], + ]), + dofs_idx, + ) + print("=== force range ===\n", franka.get_dofs_force_range()) + franka.set_dofs_armature( + np.array([ + [0.1] * len(dofs_idx), + [0.2] * len(dofs_idx), + ]), + dofs_idx, + ) + print("=== armature ===\n", franka.get_dofs_armature()) + franka.set_dofs_stiffness( + np.array([ + [0.] * len(dofs_idx), + [0.1] * len(dofs_idx), + ]), + dofs_idx, + ) + print("=== stiffness ===\n", franka.get_dofs_stiffness()) + franka.set_dofs_invweight( + np.array([ + [5.5882, 0.9693, 6.8053, 3.9007, 7.8085, 6.6139, 9.4213, 8.6984, 8.6984], + [8.6984, 8.6984, 9.4213, 6.6139, 7.8085, 3.9007, 6.8053, 0.9693, 5.5882], + ]), + dofs_idx, + ) + print("=== invweight ===\n", franka.get_dofs_invweight()) + franka.set_dofs_damping( + np.array([ + [1.] * len(dofs_idx), + [2.] * len(dofs_idx), + ]), + dofs_idx, + ) + print("=== damping ===\n", franka.get_dofs_damping()) + franka.set_links_inertial_mass( + np.array([ + [0.6298, 4.9707, 0.6469, 3.2286, 3.5879, 1.2259, 1.6666, 0.7355, 0.7300, 0.0150, 0.0150], + [0.015, 0.015, 0.73, 0.7355, 1.6666, 1.2259, 3.5879, 3.2286, 0.6469, 4.9707, 0.6298], + ]), + links_idx, + ) + print("=== links inertial mass ===\n", franka.get_links_inertial_mass()) + franka.set_links_invweight( + np.array([ + [0.0, 3.6037e-05, 0.00030664, 0.025365, 0.036351, 0.072328, 0.089559, 0.11661, 0.11288, 3.0179, 3.0179], + [3.0179, 3.0179, 0.11288, 0.11661, 0.089559, 0.072328, 0.036351, 0.025365, 0.00030664, 3.6037e-05, 0.0], + ]), + links_idx, + ) + print("=== links invweight ===\n", franka.get_links_invweight()) + + # Hard reset + for i in range(150): + if i < 50: + franka.set_dofs_position(np.array([1, 1, 0, 0, 0, 0, 0, 0.04, 0.04])[None, :].repeat(scene.n_envs, 0), dofs_idx) + elif i < 100: + franka.set_dofs_position(np.array([-1, 0.8, 1, -2, 1, 0.5, -0.5, 0.04, 0.04])[None, :].repeat(scene.n_envs, 0), dofs_idx) + else: + franka.set_dofs_position(np.array([0, 0, 0, 0, 0, 0, 0, 0, 0])[None, :].repeat(scene.n_envs, 0), dofs_idx) + + scene.step() + + # PD control + for i in range(1250): + if i == 0: + franka.control_dofs_position( + np.array([1, 1, 0, 0, 0, 0, 0, 0.04, 0.04])[None, :].repeat(scene.n_envs, 0), + dofs_idx, + ) + elif i == 250: + franka.control_dofs_position( + np.array([-1, 0.8, 1, -2, 1, 0.5, -0.5, 0.04, 0.04])[None, :].repeat(scene.n_envs, 0), + dofs_idx, + ) + elif i == 500: + franka.control_dofs_position( + np.array([0, 0, 0, 0, 0, 0, 0, 0, 0])[None, :].repeat(scene.n_envs, 0), + dofs_idx, + ) + elif i == 750: + # control first dof with velocity, and the rest with position + franka.control_dofs_position( + np.array([0, 0, 0, 0, 0, 0, 0, 0, 0])[1:][None, :].repeat(scene.n_envs, 0), + dofs_idx[1:], + ) + franka.control_dofs_velocity( + np.array([1.0, 0, 0, 0, 0, 0, 0, 0, 0])[:1][None, :].repeat(scene.n_envs, 0), + dofs_idx[:1], + ) + elif i == 1000: + franka.control_dofs_force( + np.array([0, 0, 0, 0, 0, 0, 0, 0, 0])[None, :].repeat(scene.n_envs, 0), + dofs_idx, + ) + # This is the internal control force computed based on the given control command + # If using force control, it's the same as the given control command + print("control force:", franka.get_dofs_control_force(dofs_idx)) + + scene.step() + + +if __name__ == "__main__": + main() diff --git a/genesis/engine/entities/rigid_entity/rigid_entity.py b/genesis/engine/entities/rigid_entity/rigid_entity.py index 82f6df00..967530f9 100644 --- a/genesis/engine/entities/rigid_entity/rigid_entity.py +++ b/genesis/engine/entities/rigid_entity/rigid_entity.py @@ -1567,6 +1567,14 @@ def get_links_ang(self, envs_idx=None): """ return self._solver.get_links_ang(np.arange(self.link_start, self.link_end), envs_idx) + @gs.assert_built + def get_links_inertial_mass(self, ls_idx_local=None, envs_idx=None): + return self._solver.get_links_inertial_mass(self._get_ls_idx(ls_idx_local), envs_idx) + + @gs.assert_built + def get_links_invweight(self, ls_idx_local=None, envs_idx=None): + return self._solver.get_links_invweight(self._get_ls_idx(ls_idx_local), envs_idx) + @gs.assert_built def set_pos(self, pos, zero_velocity=True, envs_idx=None): """ @@ -1700,6 +1708,18 @@ def _get_dofs_idx_local(self, dofs_idx_local=None): def _get_dofs_idx(self, dofs_idx_local=None): return self._get_dofs_idx_local(dofs_idx_local) + self._dof_start + def _get_ls_idx_local(self, ls_idx_local=None): + if ls_idx_local is None: + ls_idx_local = torch.arange(self.n_links, dtype=torch.int32, device=gs.device) + else: + ls_idx_local = torch.as_tensor(ls_idx_local, dtype=gs.tc_int) + if (ls_idx_local < 0).any() or (ls_idx_local >= self.n_links).any(): + gs.raise_exception('`ls_idx_local` exceeds valid range.') + return ls_idx_local + + def _get_ls_idx(self, ls_idx_local=None): + return self._get_ls_idx_local(ls_idx_local) + self._link_start + @gs.assert_built def set_qpos(self, qpos, qs_idx_local=None, zero_velocity=True, envs_idx=None): """ @@ -1773,6 +1793,22 @@ def set_dofs_force_range(self, lower, upper, dofs_idx_local=None, envs_idx=None) self._solver.set_dofs_force_range(lower, upper, self._get_dofs_idx(dofs_idx_local), envs_idx) + @gs.assert_built + def set_dofs_stiffness(self, stiffness, dofs_idx_local=None, envs_idx=None): + self._solver.set_dofs_stiffness(stiffness, self._get_dofs_idx(dofs_idx_local), envs_idx) + + @gs.assert_built + def set_dofs_invweight(self, invweight, dofs_idx_local=None, envs_idx=None): + self._solver.set_dofs_invweight(invweight, self._get_dofs_idx(dofs_idx_local), envs_idx) + + @gs.assert_built + def set_dofs_armature(self, armature, dofs_idx_local=None, envs_idx=None): + self._solver.set_dofs_armature(armature, self._get_dofs_idx(dofs_idx_local), envs_idx) + + @gs.assert_built + def set_dofs_damping(self, damping, dofs_idx_local=None, envs_idx=None): + self._solver.set_dofs_damping(damping, self._get_dofs_idx(dofs_idx_local), envs_idx) + @gs.assert_built def set_dofs_velocity(self, velocity, dofs_idx_local=None, envs_idx=None): """ @@ -2039,6 +2075,22 @@ def get_dofs_limit(self, dofs_idx=None, envs_idx=None): """ return self._solver.get_dofs_limit(self._get_dofs_idx(dofs_idx), envs_idx) + @gs.assert_built + def get_dofs_stiffness(self, dofs_idx_local=None, envs_idx=None): + return self._solver.get_dofs_stiffness(self._get_dofs_idx(dofs_idx_local), envs_idx) + + @gs.assert_built + def get_dofs_invweight(self, dofs_idx_local=None, envs_idx=None): + return self._solver.get_dofs_invweight(self._get_dofs_idx(dofs_idx_local), envs_idx) + + @gs.assert_built + def get_dofs_armature(self, dofs_idx_local=None, envs_idx=None): + return self._solver.get_dofs_armature(self._get_dofs_idx(dofs_idx_local), envs_idx) + + @gs.assert_built + def get_dofs_damping(self, dofs_idx_local=None, envs_idx=None): + return self._solver.get_dofs_damping(self._get_dofs_idx(dofs_idx_local), envs_idx) + @gs.assert_built def zero_all_dofs_velocity(self, envs_idx=None): """ @@ -2271,6 +2323,14 @@ def set_COM_shift(self, com_shift, link_indices, envs_idx=None): link_indices[i] += self._link_start self._solver.set_links_COM_shift(com_shift, link_indices, envs_idx) + @gs.assert_built + def set_links_inertial_mass(self, inertial_mass, ls_idx_local=None, envs_idx=None): + self._solver.set_links_inertial_mass(inertial_mass, self._get_ls_idx(ls_idx_local), envs_idx) + + @gs.assert_built + def set_links_invweight(self, invweight, ls_idx_local=None, envs_idx=None): + self._solver.set_links_invweight(invweight, self._get_ls_idx(ls_idx_local), envs_idx) + @gs.assert_built def get_mass(self): """ diff --git a/genesis/engine/solvers/rigid/rigid_solver_decomp.py b/genesis/engine/solvers/rigid/rigid_solver_decomp.py index b045f694..b20ced10 100644 --- a/genesis/engine/solvers/rigid/rigid_solver_decomp.py +++ b/genesis/engine/solvers/rigid/rigid_solver_decomp.py @@ -3356,6 +3356,56 @@ def _kernel_set_links_COM_shift( for i in ti.static(range(3)): self.links_state[links_idx[i_l_], envs_idx[i_b_]].i_pos_shift[i] = com[i_b_, i_l_, i] + def _set_links_info(self, tensor, links_idx, name, envs_idx=None): + if self._options.batch_links_info: + tensor, links_idx, envs_idx = self._validate_1D_io_variables(tensor, links_idx, envs_idx, idx_name='links_idx') + else: + tensor, links_idx = self._validate_1D_io_variables(tensor, links_idx, idx_name='links_idx', batched=False) + envs_idx = torch.empty(()) + + if name == "invweight": + self._kernel_set_links_invweight(tensor, links_idx, envs_idx) + elif name == "inertial_mass": + self._kernel_set_links_inertial_mass(tensor, links_idx, envs_idx) + else: + gs.raise_exception(f"Invalid `name` {name}.") + + def set_links_inertial_mass(self, invweight, links_idx, envs_idx=None): + self._set_links_info(invweight, links_idx, "inertial_mass", envs_idx) + + @ti.kernel + def _kernel_set_links_inertial_mass( + self, + inertial_mass: ti.types.ndarray(), + links_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + ): + ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) + if ti.static(self._options.batch_links_info): + for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]): + self.links_info[links_idx[i_l_], envs_idx[i_b_]].inertial_mass = inertial_mass[i_b_, i_l_] + else: + for i_l_ in range(links_idx.shape[0]): + self.links_info[links_idx[i_l_]].inertial_mass = inertial_mass[i_l_] + + def set_links_invweight(self, invweight, links_idx, envs_idx=None): + self._set_links_info(invweight, links_idx, "invweight", envs_idx) + + @ti.kernel + def _kernel_set_links_invweight( + self, + invweight: ti.types.ndarray(), + links_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + ): + ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) + if ti.static(self._options.batch_links_info): + for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]): + self.links_info[links_idx[i_l_], envs_idx[i_b_]].invweight = invweight[i_b_, i_l_] + else: + for i_l_ in range(links_idx.shape[0]): + self.links_info[links_idx[i_l_]].invweight = invweight[i_l_] + def set_geoms_friction_ratio(self, friction_ratio, geoms_idx, envs_idx=None): friction_ratio, geoms_idx, envs_idx = self._validate_1D_io_variables( friction_ratio, geoms_idx, envs_idx, idx_name="geoms_idx" @@ -3418,13 +3468,42 @@ def _kernel_set_global_sol_params(self, sol_params: ti.types.ndarray()): self.dofs_info[I].sol_params[0] = self._substep_dt * 2 - def set_dofs_kp(self, kp, dofs_idx, envs_idx=None): + def _set_dofs_info(self, tensor_list, dofs_idx, name, envs_idx=None): if self._options.batch_dofs_info: - kp, dofs_idx, envs_idx = self._validate_1D_io_variables(kp, dofs_idx, envs_idx) + for i, tensor in enumerate(tensor_list): + if i == (len(tensor_list) - 1): + tensor_list[i], dofs_idx, envs_idx = self._validate_1D_io_variables(tensor, dofs_idx, envs_idx) + else: + tensor_list[i], _, _ = self._validate_1D_io_variables(tensor, dofs_idx, envs_idx) else: - kp, dofs_idx = self._validate_1D_io_variables(kp, dofs_idx, batched=False) + for i, tensor in enumerate(tensor_list): + if i == (len(tensor_list) - 1): + tensor_list[i], _ = self._validate_1D_io_variables(tensor, dofs_idx, batched=False) + else: + tensor_list[i], dofs_idx = self._validate_1D_io_variables(tensor, dofs_idx, batched=False) envs_idx = torch.empty(()) - self._kernel_set_dofs_kp(kp, dofs_idx, envs_idx) + + if name == "kp": + self._kernel_set_dofs_kp(tensor_list[0], dofs_idx, envs_idx) + elif name == "kv": + self._kernel_set_dofs_kv(tensor_list[0], dofs_idx, envs_idx) + elif name == "force_range": + self._kernel_set_dofs_force_range(tensor_list[0], tensor_list[1], dofs_idx, envs_idx) + elif name == "stiffness": + self._kernel_set_dofs_stiffness(tensor_list[0], dofs_idx, envs_idx) + elif name == "invweight": + self._kernel_set_dofs_invweight(tensor_list[0], dofs_idx, envs_idx) + elif name == "armature": + self._kernel_set_dofs_armature(tensor_list[0], dofs_idx, envs_idx) + elif name == "damping": + self._kernel_set_dofs_damping(tensor_list[0], dofs_idx, envs_idx) + elif name == "limit": + self._kernel_set_dofs_limit(tensor_list[0], tensor_list[1], dofs_idx, envs_idx) + else: + gs.raise_exception(f"Invalid `name` {name}.") + + def set_dofs_kp(self, kp, dofs_idx, envs_idx=None): + self._set_dofs_info([kp], dofs_idx, "kp", envs_idx) @ti.kernel def _kernel_set_dofs_kp( @@ -3442,12 +3521,7 @@ def _kernel_set_dofs_kp( self.dofs_info[dofs_idx[i_d_]].kp = kp[i_d_] def set_dofs_kv(self, kv, dofs_idx, envs_idx=None): - if self._options.batch_dofs_info: - kv, dofs_idx, envs_idx = self._validate_1D_io_variables(kv, dofs_idx, envs_idx) - else: - kv, dofs_idx = self._validate_1D_io_variables(kv, dofs_idx, batched=False) - envs_idx = torch.empty(()) - self._kernel_set_dofs_kv(kv, dofs_idx, envs_idx) + self._set_dofs_info([kv], dofs_idx, "kv", envs_idx) @ti.kernel def _kernel_set_dofs_kv( @@ -3465,18 +3539,7 @@ def _kernel_set_dofs_kv( self.dofs_info[dofs_idx[i_d_]].kv = kv[i_d_] def set_dofs_force_range(self, lower, upper, dofs_idx, envs_idx=None): - if self._options.batch_dofs_info: - lower, _, _ = self._validate_1D_io_variables(lower, dofs_idx, envs_idx) - upper, dofs_idx, envs_idx = self._validate_1D_io_variables(upper, dofs_idx, envs_idx) - else: - lower, _ = self._validate_1D_io_variables(lower, dofs_idx, batched=False) - upper, dofs_idx = self._validate_1D_io_variables(upper, dofs_idx, batched=False) - envs_idx = torch.empty(()) - - if (lower > upper).any(): - gs.raise_exception("`lower` should be less than or equal to `upper`.") - - self._kernel_set_dofs_force_range(lower, upper, dofs_idx, envs_idx) + self._set_dofs_info([lower, upper], dofs_idx, "force_range", envs_idx) @ti.kernel def _kernel_set_dofs_force_range( @@ -3496,6 +3559,99 @@ def _kernel_set_dofs_force_range( self.dofs_info[dofs_idx[i_d_]].force_range[0] = lower[i_d_] self.dofs_info[dofs_idx[i_d_]].force_range[1] = upper[i_d_] + def set_dofs_stiffness(self, stiffness, dofs_idx, envs_idx=None): + self._set_dofs_info([stiffness], dofs_idx, "stiffness", envs_idx) + + @ti.kernel + def _kernel_set_dofs_stiffness( + self, + stiffness: ti.types.ndarray(), + dofs_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + ): + ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].stiffness = stiffness[i_b_, i_d_] + else: + for i_d_ in range(dofs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_]].stiffness = stiffness[i_d_] + + def set_dofs_invweight(self, invweight, dofs_idx, envs_idx=None): + self._set_dofs_info([invweight], dofs_idx, "invweight", envs_idx) + + @ti.kernel + def _kernel_set_dofs_invweight( + self, + invweight: ti.types.ndarray(), + dofs_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + ): + ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].invweight = invweight[i_b_, i_d_] + else: + for i_d_ in range(dofs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_]].invweight = invweight[i_d_] + + def set_dofs_armature(self, armature, dofs_idx, envs_idx=None): + self._set_dofs_info([armature], dofs_idx, "armature", envs_idx) + + @ti.kernel + def _kernel_set_dofs_armature( + self, + armature: ti.types.ndarray(), + dofs_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + ): + ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].armature = armature[i_b_, i_d_] + else: + for i_d_ in range(dofs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_]].armature = armature[i_d_] + + def set_dofs_damping(self, damping, dofs_idx, envs_idx=None): + self._set_dofs_info([damping], dofs_idx, "damping", envs_idx) + + @ti.kernel + def _kernel_set_dofs_damping( + self, + damping: ti.types.ndarray(), + dofs_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + ): + ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].damping = damping[i_b_, i_d_] + else: + for i_d_ in range(dofs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_]].damping = damping[i_d_] + + def set_dofs_limit(self, lower, upper, dofs_idx, envs_idx=None): + self._set_dofs_info([lower, upper], dofs_idx, "limit", envs_idx) + + @ti.kernel + def _kernel_set_dofs_limit( + self, + lower: ti.types.ndarray(), + upper: ti.types.ndarray(), + dofs_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + ): + ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].limit[0] = lower[i_b_, i_d_] + self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].limit[1] = upper[i_b_, i_d_] + else: + for i_d_ in range(dofs_idx.shape[0]): + self.dofs_info[dofs_idx[i_d_]].limit[0] = lower[i_d_] + self.dofs_info[dofs_idx[i_d_]].limit[1] = upper[i_d_] + def set_dofs_velocity(self, velocity, dofs_idx, envs_idx=None): velocity, dofs_idx, envs_idx = self._validate_1D_io_variables(velocity, dofs_idx, envs_idx) self._kernel_set_dofs_velocity( @@ -3770,6 +3926,58 @@ def _kernel_get_links_COM_shift( for i in ti.static(range(3)): tensor[i_b_, i_l_, i] = self.links_state[links_idx[i_l_], envs_idx[i_b_]].i_pos_shift[i] + def _get_links_info(self, links_idx, name, envs_idx=None): + if self._options.batch_links_info: + tensor, links_idx, envs_idx = self._validate_1D_io_variables(None, links_idx, envs_idx, idx_name="links_idx") + else: + tensor, links_idx = self._validate_1D_io_variables(None, links_idx, idx_name="links_idx", batched=False) + envs_idx = torch.empty(()) + + if name == "invweight": + self._kernel_get_links_invweight(tensor, links_idx, envs_idx) + return tensor + elif name == "inertial_mass": + self._kernel_get_links_inertial_mass(tensor, links_idx, envs_idx) + return tensor + else: + gs.raise_exception(f"Invalid `name` {name}.") + + def get_links_inertial_mass(self, links_idx, envs_idx=None): + return self._get_links_info(links_idx, "inertial_mass", envs_idx) + + @ti.kernel + def _kernel_get_links_inertial_mass( + self, + tensor: ti.types.ndarray(), + links_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + ): + ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) + if ti.static(self._options.batch_links_info): + for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]): + tensor[i_b_, i_l_] = self.links_info[links_idx[i_l_], envs_idx[i_b_]].inertial_mass + else: + for i_l_ in range(links_idx.shape[0]): + tensor[i_l_] = self.links_info[links_idx[i_l_]].inertial_mass + + def get_links_invweight(self, links_idx, envs_idx=None): + return self._get_links_info(links_idx, "invweight", envs_idx) + + @ti.kernel + def _kernel_get_links_invweight( + self, + tensor: ti.types.ndarray(), + links_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + ): + ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) + if ti.static(self._options.batch_links_info): + for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]): + tensor[i_b_, i_l_] = self.links_info[links_idx[i_l_], envs_idx[i_b_]].invweight + else: + for i_l_ in range(links_idx.shape[0]): + tensor[i_l_] = self.links_info[links_idx[i_l_]].invweight + def get_geoms_friction_ratio(self, geoms_idx, envs_idx=None): tensor, geoms_idx, envs_idx = self._validate_1D_io_variables(None, geoms_idx, envs_idx, idx_name="geoms_idx") @@ -3935,6 +4143,18 @@ def get_dofs_force_range(self, dofs_idx, envs_idx=None): def get_dofs_limit(self, dofs_idx, envs_idx=None): return self._get_dofs_info(dofs_idx, "limit", envs_idx) + def get_dofs_stiffness(self, dofs_idx, envs_idx=None): + return self._get_dofs_info(dofs_idx, "stiffness", envs_idx) + + def get_dofs_invweight(self, dofs_idx, envs_idx=None): + return self._get_dofs_info(dofs_idx, "invweight", envs_idx) + + def get_dofs_armature(self, dofs_idx, envs_idx=None): + return self._get_dofs_info(dofs_idx, "armature", envs_idx) + + def get_dofs_damping(self, dofs_idx, envs_idx=None): + return self._get_dofs_info(dofs_idx, "damping", envs_idx) + def _get_dofs_info(self, dofs_idx, name, envs_idx=None): if self._options.batch_dofs_info: tensor, dofs_idx, envs_idx = self._validate_1D_io_variables(None, dofs_idx, envs_idx) @@ -3962,6 +4182,22 @@ def _get_dofs_info(self, dofs_idx, name, envs_idx=None): self._kernel_get_dofs_limit(lower, upper, dofs_idx, envs_idx) return lower, upper + elif name == "stiffness": + self._kernel_get_dofs_stiffness(tensor, dofs_idx, envs_idx) + return tensor + + elif name == "invweight": + self._kernel_get_dofs_invweight(tensor, dofs_idx, envs_idx) + return tensor + + elif name == "armature": + self._kernel_get_dofs_armature(tensor, dofs_idx, envs_idx) + return tensor + + elif name == "damping": + self._kernel_get_dofs_damping(tensor, dofs_idx, envs_idx) + return tensor + else: gs.raise_exception() @@ -4031,6 +4267,66 @@ def _kernel_get_dofs_limit( lower[i_d_] = self.dofs_info[dofs_idx[i_d_]].limit[0] upper[i_d_] = self.dofs_info[dofs_idx[i_d_]].limit[1] + @ti.kernel + def _kernel_get_dofs_stiffness( + self, + tensor: ti.types.ndarray(), + dofs_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + ): + ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + tensor[i_b_, i_d_] = self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].stiffness + else: + for i_d_ in range(dofs_idx.shape[0]): + tensor[i_d_] = self.dofs_info[dofs_idx[i_d_]].stiffness + + @ti.kernel + def _kernel_get_dofs_invweight( + self, + tensor: ti.types.ndarray(), + dofs_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + ): + ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + tensor[i_b_, i_d_] = self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].invweight + else: + for i_d_ in range(dofs_idx.shape[0]): + tensor[i_d_] = self.dofs_info[dofs_idx[i_d_]].invweight + + @ti.kernel + def _kernel_get_dofs_armature( + self, + tensor: ti.types.ndarray(), + dofs_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + ): + ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + tensor[i_b_, i_d_] = self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].armature + else: + for i_d_ in range(dofs_idx.shape[0]): + tensor[i_d_] = self.dofs_info[dofs_idx[i_d_]].armature + + @ti.kernel + def _kernel_get_dofs_damping( + self, + tensor: ti.types.ndarray(), + dofs_idx: ti.types.ndarray(), + envs_idx: ti.types.ndarray(), + ): + ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) + if ti.static(self._options.batch_dofs_info): + for i_d_, i_b_ in ti.ndrange(dofs_idx.shape[0], envs_idx.shape[0]): + tensor[i_b_, i_d_] = self.dofs_info[dofs_idx[i_d_], envs_idx[i_b_]].damping + else: + for i_d_ in range(dofs_idx.shape[0]): + tensor[i_d_] = self.dofs_info[dofs_idx[i_d_]].damping + @ti.kernel def _kernel_set_drone_rpm( self, From 22590238b019b4552b60e5eb6e53e7a9d951b17f Mon Sep 17 00:00:00 2001 From: Johnson Wang Date: Sun, 22 Dec 2024 09:47:12 -0500 Subject: [PATCH 06/35] fix formatting --- examples/rigid/set_phys_attr.py | 130 +++++++++++------- .../entities/rigid_entity/rigid_entity.py | 8 +- .../solvers/rigid/constraint_solver_decomp.py | 10 +- .../rigid/constraint_solver_decomp_island.py | 4 +- .../solvers/rigid/rigid_solver_decomp.py | 16 ++- 5 files changed, 103 insertions(+), 65 deletions(-) diff --git a/examples/rigid/set_phys_attr.py b/examples/rigid/set_phys_attr.py index 6979c5ec..8dc0bf56 100644 --- a/examples/rigid/set_phys_attr.py +++ b/examples/rigid/set_phys_attr.py @@ -43,7 +43,7 @@ def main(): gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"), ) ########################## build ########################## - scene.build(n_envs=2) # test with 2 different environments + scene.build(n_envs=2) # test with 2 different environments jnt_names = [ "joint1", @@ -60,93 +60,113 @@ def main(): lnk_names = [ "link0", - 'link1', - 'link2', - 'link3', - 'link4', - 'link5', - 'link6', - 'link7', - 'hand', - 'left_finger', - 'right_finger', + "link1", + "link2", + "link3", + "link4", + "link5", + "link6", + "link7", + "hand", + "left_finger", + "right_finger", ] links_idx = [franka.get_link(name).idx_local for name in lnk_names] # Optional: set control gains franka.set_dofs_kp( - np.array([ - [4500, 4500, 3500, 3500, 2000, 2000, 2000, 100, 100], - [100, 100, 2000, 2000, 2000, 3500, 3500, 4500, 4500], - ]), + np.array( + [ + [4500, 4500, 3500, 3500, 2000, 2000, 2000, 100, 100], + [100, 100, 2000, 2000, 2000, 3500, 3500, 4500, 4500], + ] + ), dofs_idx, ) print("=== kp ===\n", franka.get_dofs_kp()) franka.set_dofs_kv( - np.array([ - [450, 450, 350, 350, 200, 200, 200, 10, 10], - [10, 10, 200, 200, 200, 350, 350, 450, 450], - ]), + np.array( + [ + [450, 450, 350, 350, 200, 200, 200, 10, 10], + [10, 10, 200, 200, 200, 350, 350, 450, 450], + ] + ), dofs_idx, ) print("=== kv ===\n", franka.get_dofs_kv()) franka.set_dofs_force_range( - np.array([ - [-87, -87, -87, -87, -12, -12, -12, -100, -100], - [-120, -100, -12, -12, -12, -87, -87, -87, -87], - ]), - np.array([ - [87, 87, 87, 87, 12, 12, 12, 100, 100], - [100, 100, 12, 12, 12, 87, 87, 87, 87], - ]), + np.array( + [ + [-87, -87, -87, -87, -12, -12, -12, -100, -100], + [-120, -100, -12, -12, -12, -87, -87, -87, -87], + ] + ), + np.array( + [ + [87, 87, 87, 87, 12, 12, 12, 100, 100], + [100, 100, 12, 12, 12, 87, 87, 87, 87], + ] + ), dofs_idx, ) print("=== force range ===\n", franka.get_dofs_force_range()) franka.set_dofs_armature( - np.array([ - [0.1] * len(dofs_idx), - [0.2] * len(dofs_idx), - ]), + np.array( + [ + [0.1] * len(dofs_idx), + [0.2] * len(dofs_idx), + ] + ), dofs_idx, ) print("=== armature ===\n", franka.get_dofs_armature()) franka.set_dofs_stiffness( - np.array([ - [0.] * len(dofs_idx), - [0.1] * len(dofs_idx), - ]), + np.array( + [ + [0.0] * len(dofs_idx), + [0.1] * len(dofs_idx), + ] + ), dofs_idx, ) print("=== stiffness ===\n", franka.get_dofs_stiffness()) franka.set_dofs_invweight( - np.array([ - [5.5882, 0.9693, 6.8053, 3.9007, 7.8085, 6.6139, 9.4213, 8.6984, 8.6984], - [8.6984, 8.6984, 9.4213, 6.6139, 7.8085, 3.9007, 6.8053, 0.9693, 5.5882], - ]), + np.array( + [ + [5.5882, 0.9693, 6.8053, 3.9007, 7.8085, 6.6139, 9.4213, 8.6984, 8.6984], + [8.6984, 8.6984, 9.4213, 6.6139, 7.8085, 3.9007, 6.8053, 0.9693, 5.5882], + ] + ), dofs_idx, ) print("=== invweight ===\n", franka.get_dofs_invweight()) franka.set_dofs_damping( - np.array([ - [1.] * len(dofs_idx), - [2.] * len(dofs_idx), - ]), + np.array( + [ + [1.0] * len(dofs_idx), + [2.0] * len(dofs_idx), + ] + ), dofs_idx, ) print("=== damping ===\n", franka.get_dofs_damping()) franka.set_links_inertial_mass( - np.array([ - [0.6298, 4.9707, 0.6469, 3.2286, 3.5879, 1.2259, 1.6666, 0.7355, 0.7300, 0.0150, 0.0150], - [0.015, 0.015, 0.73, 0.7355, 1.6666, 1.2259, 3.5879, 3.2286, 0.6469, 4.9707, 0.6298], - ]), + np.array( + [ + [0.6298, 4.9707, 0.6469, 3.2286, 3.5879, 1.2259, 1.6666, 0.7355, 0.7300, 0.0150, 0.0150], + [0.015, 0.015, 0.73, 0.7355, 1.6666, 1.2259, 3.5879, 3.2286, 0.6469, 4.9707, 0.6298], + ] + ), links_idx, ) print("=== links inertial mass ===\n", franka.get_links_inertial_mass()) franka.set_links_invweight( - np.array([ - [0.0, 3.6037e-05, 0.00030664, 0.025365, 0.036351, 0.072328, 0.089559, 0.11661, 0.11288, 3.0179, 3.0179], - [3.0179, 3.0179, 0.11288, 0.11661, 0.089559, 0.072328, 0.036351, 0.025365, 0.00030664, 3.6037e-05, 0.0], - ]), + np.array( + [ + [0.0, 3.6037e-05, 0.00030664, 0.025365, 0.036351, 0.072328, 0.089559, 0.11661, 0.11288, 3.0179, 3.0179], + [3.0179, 3.0179, 0.11288, 0.11661, 0.089559, 0.072328, 0.036351, 0.025365, 0.00030664, 3.6037e-05, 0.0], + ] + ), links_idx, ) print("=== links invweight ===\n", franka.get_links_invweight()) @@ -154,9 +174,13 @@ def main(): # Hard reset for i in range(150): if i < 50: - franka.set_dofs_position(np.array([1, 1, 0, 0, 0, 0, 0, 0.04, 0.04])[None, :].repeat(scene.n_envs, 0), dofs_idx) + franka.set_dofs_position( + np.array([1, 1, 0, 0, 0, 0, 0, 0.04, 0.04])[None, :].repeat(scene.n_envs, 0), dofs_idx + ) elif i < 100: - franka.set_dofs_position(np.array([-1, 0.8, 1, -2, 1, 0.5, -0.5, 0.04, 0.04])[None, :].repeat(scene.n_envs, 0), dofs_idx) + franka.set_dofs_position( + np.array([-1, 0.8, 1, -2, 1, 0.5, -0.5, 0.04, 0.04])[None, :].repeat(scene.n_envs, 0), dofs_idx + ) else: franka.set_dofs_position(np.array([0, 0, 0, 0, 0, 0, 0, 0, 0])[None, :].repeat(scene.n_envs, 0), dofs_idx) diff --git a/genesis/engine/entities/rigid_entity/rigid_entity.py b/genesis/engine/entities/rigid_entity/rigid_entity.py index 967530f9..955af57c 100644 --- a/genesis/engine/entities/rigid_entity/rigid_entity.py +++ b/genesis/engine/entities/rigid_entity/rigid_entity.py @@ -1162,7 +1162,11 @@ def _kernel_inverse_kinematics( for i_l in range(self.link_start, self.link_end): I_l = [i_l, i_b] if ti.static(self.solver._options.batch_links_info) else i_l l_info = self._solver.links_info[I_l] - I_dof_start = [l_info.dof_start, i_b] if ti.static(self.solver._options.batch_dofs_info) else l_info.dof_start + I_dof_start = ( + [l_info.dof_start, i_b] + if ti.static(self.solver._options.batch_dofs_info) + else l_info.dof_start + ) dof_info = self._solver.dofs_info[I_dof_start] q_start = l_info.q_start @@ -1714,7 +1718,7 @@ def _get_ls_idx_local(self, ls_idx_local=None): else: ls_idx_local = torch.as_tensor(ls_idx_local, dtype=gs.tc_int) if (ls_idx_local < 0).any() or (ls_idx_local >= self.n_links).any(): - gs.raise_exception('`ls_idx_local` exceeds valid range.') + gs.raise_exception("`ls_idx_local` exceeds valid range.") return ls_idx_local def _get_ls_idx(self, ls_idx_local=None): diff --git a/genesis/engine/solvers/rigid/constraint_solver_decomp.py b/genesis/engine/solvers/rigid/constraint_solver_decomp.py index 43c0cab1..d3592b92 100644 --- a/genesis/engine/solvers/rigid/constraint_solver_decomp.py +++ b/genesis/engine/solvers/rigid/constraint_solver_decomp.py @@ -112,9 +112,9 @@ def add_collision_constraints(self): d1, d2 = gu.orthogonals(impact.normal) - t = self._solver.links_info[link_a_maybe_batch].invweight + self._solver.links_info[link_b_maybe_batch].invweight * ( - link_b > -1 - ) + t = self._solver.links_info[link_a_maybe_batch].invweight + self._solver.links_info[ + link_b_maybe_batch + ].invweight * (link_b > -1) for i in range(4): n = -d1 * f - impact.normal if i == 1: @@ -145,7 +145,9 @@ def add_collision_constraints(self): link = link_b while link > -1: - link_maybe_batch = [link, i_b] if ti.static(self._solver._options.batch_links_info) else link + link_maybe_batch = ( + [link, i_b] if ti.static(self._solver._options.batch_links_info) else link + ) # reverse order to make sure dofs in each row of self.jac_relevant_dofs is strictly descending for i_d_ in range(self._solver.links_info[link_maybe_batch].n_dofs): diff --git a/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py b/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py index 534686fa..6f69dd2d 100644 --- a/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py +++ b/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py @@ -138,7 +138,9 @@ def add_collision_constraints(self, island, i_b): d1, d2 = gu.orthogonals(impact.normal) - t = self._solver.links_info[link_a_maybe_batch].invweight + self._solver.links_info[link_b_maybe_batch].invweight * (link_b > -1) + t = self._solver.links_info[link_a_maybe_batch].invweight + self._solver.links_info[ + link_b_maybe_batch + ].invweight * (link_b > -1) for i in range(4): n = -d1 * f - impact.normal diff --git a/genesis/engine/solvers/rigid/rigid_solver_decomp.py b/genesis/engine/solvers/rigid/rigid_solver_decomp.py index b20ced10..7f68c632 100644 --- a/genesis/engine/solvers/rigid/rigid_solver_decomp.py +++ b/genesis/engine/solvers/rigid/rigid_solver_decomp.py @@ -346,7 +346,7 @@ def _kernel_init_dof_fields( ): ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for I in ti.grouped(self.dofs_info): - i = I[0] # batching (if any) will be the second dim + i = I[0] # batching (if any) will be the second dim for j in ti.static(range(3)): self.dofs_info[I].motion_ang[j] = dofs_motion_ang[i, j] @@ -2499,7 +2499,9 @@ def _func_torque_and_passive_force(self): q_end = l_info.q_end for j_d in range(q_end - q_start): - I_d = [dof_start + j_d, i_b] if ti.static(self._options.batch_dofs_info) else dof_start + j_d + I_d = ( + [dof_start + j_d, i_b] if ti.static(self._options.batch_dofs_info) else dof_start + j_d + ) self.dofs_state[dof_start + j_d, i_b].qf_passive = ( -self.qpos[q_start + j_d, i_b] * self.dofs_info[I_d].stiffness ) @@ -3358,9 +3360,11 @@ def _kernel_set_links_COM_shift( def _set_links_info(self, tensor, links_idx, name, envs_idx=None): if self._options.batch_links_info: - tensor, links_idx, envs_idx = self._validate_1D_io_variables(tensor, links_idx, envs_idx, idx_name='links_idx') + tensor, links_idx, envs_idx = self._validate_1D_io_variables( + tensor, links_idx, envs_idx, idx_name="links_idx" + ) else: - tensor, links_idx = self._validate_1D_io_variables(tensor, links_idx, idx_name='links_idx', batched=False) + tensor, links_idx = self._validate_1D_io_variables(tensor, links_idx, idx_name="links_idx", batched=False) envs_idx = torch.empty(()) if name == "invweight": @@ -3928,7 +3932,9 @@ def _kernel_get_links_COM_shift( def _get_links_info(self, links_idx, name, envs_idx=None): if self._options.batch_links_info: - tensor, links_idx, envs_idx = self._validate_1D_io_variables(None, links_idx, envs_idx, idx_name="links_idx") + tensor, links_idx, envs_idx = self._validate_1D_io_variables( + None, links_idx, envs_idx, idx_name="links_idx" + ) else: tensor, links_idx = self._validate_1D_io_variables(None, links_idx, idx_name="links_idx", batched=False) envs_idx = torch.empty(()) From e008185cef81e99cd83f8dae7425380299775d60 Mon Sep 17 00:00:00 2001 From: Johnson Wang Date: Sun, 22 Dec 2024 10:13:46 -0500 Subject: [PATCH 07/35] fix minor --- .../engine/solvers/rigid/constraint_solver_decomp_island.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py b/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py index 6f69dd2d..4a7d2257 100644 --- a/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py +++ b/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py @@ -224,13 +224,13 @@ def add_joint_limit_constraints(self, island, i_b): e_info = self.entities_info[i_e] for i_l in range(e_info.link_start, e_info.link_end): - I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + I_l = [i_l, i_b] if ti.static(self._solver._options.batch_links_info) else i_l l_info = self._solver.links_info[I_l] if l_info.joint_type == gs.JOINT_TYPE.REVOLUTE or l_info.joint_type == gs.JOINT_TYPE.PRISMATIC: i_q = l_info.q_start i_d = l_info.dof_start - I_d = [i_d, i_b] if ti.static(self._options.batch_dofs_info) else i_d + I_d = [i_d, i_b] if ti.static(self._solver._options.batch_dofs_info) else i_d pos_min = self._solver.qpos[i_q, i_b] - self._solver.dofs_info[I_d].limit[0] pos_max = self._solver.dofs_info[I_d].limit[1] - self._solver.qpos[i_q, i_b] pos = min(min(pos_min, pos_max), 0) From 87c53a8f1a52bc8d10aa6c9a22312ee3b76c7206 Mon Sep 17 00:00:00 2001 From: Johnson Wang Date: Sun, 22 Dec 2024 13:29:06 -0500 Subject: [PATCH 08/35] test all examples; fix empty dofs; fix minor --- .../solvers/rigid/rigid_solver_decomp.py | 32 ++++++++++--------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/genesis/engine/solvers/rigid/rigid_solver_decomp.py b/genesis/engine/solvers/rigid/rigid_solver_decomp.py index 7f68c632..806b8b9d 100644 --- a/genesis/engine/solvers/rigid/rigid_solver_decomp.py +++ b/genesis/engine/solvers/rigid/rigid_solver_decomp.py @@ -312,19 +312,21 @@ def _init_dof_fields(self): ) joints = self.joints - self._kernel_init_dof_fields( - dofs_motion_ang=np.concatenate([joint.dofs_motion_ang for joint in joints], dtype=gs.np_float), - dofs_motion_vel=np.concatenate([joint.dofs_motion_vel for joint in joints], dtype=gs.np_float), - dofs_limit=np.concatenate([joint.dofs_limit for joint in joints], dtype=gs.np_float), - dofs_invweight=np.concatenate([joint.dofs_invweight for joint in joints], dtype=gs.np_float), - dofs_stiffness=np.concatenate([joint.dofs_stiffness for joint in joints], dtype=gs.np_float), - dofs_sol_params=np.concatenate([joint.dofs_sol_params for joint in joints], dtype=gs.np_float), - dofs_damping=np.concatenate([joint.dofs_damping for joint in joints], dtype=gs.np_float), - dofs_armature=np.concatenate([joint.dofs_armature for joint in joints], dtype=gs.np_float), - dofs_kp=np.concatenate([joint.dofs_kp for joint in joints], dtype=gs.np_float), - dofs_kv=np.concatenate([joint.dofs_kv for joint in joints], dtype=gs.np_float), - dofs_force_range=np.concatenate([joint.dofs_force_range for joint in joints], dtype=gs.np_float), - ) + is_nonempty = np.concatenate([joint.dofs_motion_ang for joint in joints], dtype=gs.np_float).shape[0] > 0 + if is_nonempty: # handle the case where there is a link with no dofs -- otherwise may cause invalid memory + self._kernel_init_dof_fields( + dofs_motion_ang=np.concatenate([joint.dofs_motion_ang for joint in joints], dtype=gs.np_float), + dofs_motion_vel=np.concatenate([joint.dofs_motion_vel for joint in joints], dtype=gs.np_float), + dofs_limit=np.concatenate([joint.dofs_limit for joint in joints], dtype=gs.np_float), + dofs_invweight=np.concatenate([joint.dofs_invweight for joint in joints], dtype=gs.np_float), + dofs_stiffness=np.concatenate([joint.dofs_stiffness for joint in joints], dtype=gs.np_float), + dofs_sol_params=np.concatenate([joint.dofs_sol_params for joint in joints], dtype=gs.np_float), + dofs_damping=np.concatenate([joint.dofs_damping for joint in joints], dtype=gs.np_float), + dofs_armature=np.concatenate([joint.dofs_armature for joint in joints], dtype=gs.np_float), + dofs_kp=np.concatenate([joint.dofs_kp for joint in joints], dtype=gs.np_float), + dofs_kv=np.concatenate([joint.dofs_kv for joint in joints], dtype=gs.np_float), + dofs_force_range=np.concatenate([joint.dofs_force_range for joint in joints], dtype=gs.np_float), + ) # just in case self.dofs_state.force.fill(0) @@ -3281,7 +3283,7 @@ def _kernel_set_links_pos( ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]): i_l = links_idx[i_l_] - I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + I_l = [i_l, i_b_] if ti.static(self._options.batch_links_info) else i_l if self.links_info[I_l].is_fixed: # change links_state directly as the link's pose is not contained in qpos for i in ti.static(range(3)): self.links_state[i_l, envs_idx[i_b_]].pos[i] = pos[i_b_, i_l_, i] @@ -3315,7 +3317,7 @@ def _kernel_set_links_quat( ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL) for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]): i_l = links_idx[i_l_] - I_l = [i_l, i_b] if ti.static(self._options.batch_links_info) else i_l + I_l = [i_l, i_b_] if ti.static(self._options.batch_links_info) else i_l if self.links_info[I_l].is_fixed: # change links_state directly as the link's pose is not contained in qpos for i in ti.static(range(4)): self.links_state[i_l, envs_idx[i_b_]].quat[i] = quat[i_b_, i_l_, i] From 58895b693b374894ac57873cef236d6371f87365 Mon Sep 17 00:00:00 2001 From: Johnson Wang Date: Sun, 22 Dec 2024 13:31:32 -0500 Subject: [PATCH 09/35] fix formatting --- genesis/engine/solvers/rigid/rigid_solver_decomp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/genesis/engine/solvers/rigid/rigid_solver_decomp.py b/genesis/engine/solvers/rigid/rigid_solver_decomp.py index 806b8b9d..b6dd3a17 100644 --- a/genesis/engine/solvers/rigid/rigid_solver_decomp.py +++ b/genesis/engine/solvers/rigid/rigid_solver_decomp.py @@ -313,7 +313,7 @@ def _init_dof_fields(self): joints = self.joints is_nonempty = np.concatenate([joint.dofs_motion_ang for joint in joints], dtype=gs.np_float).shape[0] > 0 - if is_nonempty: # handle the case where there is a link with no dofs -- otherwise may cause invalid memory + if is_nonempty: # handle the case where there is a link with no dofs -- otherwise may cause invalid memory self._kernel_init_dof_fields( dofs_motion_ang=np.concatenate([joint.dofs_motion_ang for joint in joints], dtype=gs.np_float), dofs_motion_vel=np.concatenate([joint.dofs_motion_vel for joint in joints], dtype=gs.np_float), From a843ad0054f08a9a6c3ed069e8924d10edacd41d Mon Sep 17 00:00:00 2001 From: Johnson Wang Date: Mon, 23 Dec 2024 02:12:05 -0500 Subject: [PATCH 10/35] add particle free (different from active) --- genesis/engine/entities/mpm_entity.py | 21 +++++++++++++++++++++ genesis/engine/solvers/mpm_solver.py | 27 +++++++++++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/genesis/engine/entities/mpm_entity.py b/genesis/engine/entities/mpm_entity.py index 104b4957..d9a8a4ce 100644 --- a/genesis/engine/entities/mpm_entity.py +++ b/genesis/engine/entities/mpm_entity.py @@ -185,6 +185,27 @@ def set_muscle(self, muscle_group=None, muscle_direction=None): self.set_muscle_direction(muscle_direction) + def set_free(self, free): + self._assert_active() + + self.solver._kernel_set_free( + self._particle_start, + self._n_particles, + free, + ) + + def get_free(self): + self._assert_active() + + free = gs.zeros((self._n_particles,), dtype=int, requires_grad=False, scene=self._scene) + self.solver._kernel_get_free( + self._particle_start, + self._n_particles, + free, + ) + + return free + @ti.kernel def clear_grad(self, f: ti.i32): for i in range(self.n_particles): diff --git a/genesis/engine/solvers/mpm_solver.py b/genesis/engine/solvers/mpm_solver.py index 6356a76b..f7c9028c 100644 --- a/genesis/engine/solvers/mpm_solver.py +++ b/genesis/engine/solvers/mpm_solver.py @@ -93,6 +93,7 @@ def init_particle_fields(self): mat_idx=gs.ti_int, mass=gs.ti_float, default_Jp=gs.ti_float, + free=gs.ti_int, # for muscle muscle_group=gs.ti_int, muscle_direction=gs.ti_vec3, @@ -384,6 +385,9 @@ def p2g(self, f: ti.i32): ) self.grid[f, base - self._grid_offset + offset].mass += weight * self.particles_info[i].mass + if self.particles_info[i].free == 0: # non-free particles behave as boundary conditions + self.grid[f, base - self._grid_offset + offset].vel_in = ti.Vector.zero(gs.ti_float, 3) + @ti.kernel def g2p(self, f: ti.i32): for i in range(self._n_particles): @@ -658,6 +662,7 @@ def _kernel_add_particles( self.particles_info[i_global].mat_idx = mat_idx self.particles_info[i_global].default_Jp = mat_default_Jp self.particles_info[i_global].mass = self._p_vol * mat_rho + self.particles_info[i_global].free = 1 self.particles_info[i_global].muscle_group = 0 self.particles_info[i_global].muscle_direction = ti.Vector([0.0, 0.0, 1.0], dt=gs.ti_float) @@ -816,6 +821,28 @@ def _kernel_set_muscle_direction( for j in ti.static(range(3)): self.particles_info[i_global].muscle_direction[j] = muscle_direction[i, j] + @ti.kernel + def _kernel_set_free( + self, + particle_start: ti.i32, + n_particles: ti.i32, + free: ti.types.ndarray(), + ): + for i in range(n_particles): + i_global = i + particle_start + self.particles_info[i_global].free = free[i] + + @ti.kernel + def _kernel_get_free( + self, + particle_start: ti.i32, + n_particles: ti.i32, + free: ti.types.ndarray(), + ): + for i in range(n_particles): + i_global = i + particle_start + free[i] = self.particles_info[i_global].free + @ti.kernel def _kernel_get_state( self, From c054cf31ca20e1d9236b60628b0fb2e74ba6dce4 Mon Sep 17 00:00:00 2001 From: Johnson Wang Date: Mon, 23 Dec 2024 02:14:51 -0500 Subject: [PATCH 11/35] fix formatting --- genesis/engine/solvers/mpm_solver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/genesis/engine/solvers/mpm_solver.py b/genesis/engine/solvers/mpm_solver.py index f7c9028c..8d490022 100644 --- a/genesis/engine/solvers/mpm_solver.py +++ b/genesis/engine/solvers/mpm_solver.py @@ -385,7 +385,7 @@ def p2g(self, f: ti.i32): ) self.grid[f, base - self._grid_offset + offset].mass += weight * self.particles_info[i].mass - if self.particles_info[i].free == 0: # non-free particles behave as boundary conditions + if self.particles_info[i].free == 0: # non-free particles behave as boundary conditions self.grid[f, base - self._grid_offset + offset].vel_in = ti.Vector.zero(gs.ti_float, 3) @ti.kernel From 03e21348212fb0bbbc578226bc009f8f1e4baa99 Mon Sep 17 00:00:00 2001 From: Pierre RAFFALLI <132000863+pierridotite@users.noreply.github.com> Date: Tue, 24 Dec 2024 17:44:41 +0100 Subject: [PATCH 12/35] Update README.md (#241) [MISC] Update README.md --- README.md | 119 ++++++++++++++++++++++++++++++++++-------------------- 1 file changed, 75 insertions(+), 44 deletions(-) diff --git a/README.md b/README.md index 25c5c5a2..80c2d087 100644 --- a/README.md +++ b/README.md @@ -10,42 +10,51 @@ [![README in English](https://img.shields.io/badge/English-d9d9d9)](./README.md) [![简体中文版自述文件](https://img.shields.io/badge/简体中文-d9d9d9)](./README_CN.md) -# What is Genesis? +# Genesis -Genesis is a physics platform designed for general purpose *Robotics/Embodied AI/Physical AI* applications. It is simultaneously multiple things: +## Table of Contents + +1. [What is Genesis?](#what-is-genesis) +2. [Key Features](#key-features) +3. [Quick Installation](#quick-installation) +4. [Documentation](#documentation) +5. [Example Usage](#example-usage) +6. [Contributing to Genesis](#contributing-to-genesis) +7. [Support](#support) +8. [License and Acknowledgments](#license-and-acknowledgments) +9. [Associated Papers](#associated-papers) +10. [Citation](#citation) + +## What is Genesis? + +Genesis is a physics platform designed for general-purpose *Robotics/Embodied AI/Physical AI* applications. It is simultaneously multiple things: 1. A **universal physics engine** re-built from the ground up, capable of simulating a wide range of materials and physical phenomena. 2. A **lightweight**, **ultra-fast**, **pythonic**, and **user-friendly** robotics simulation platform. 3. A powerful and fast **photo-realistic rendering system**. 4. A **generative data engine** that transforms user-prompted natural language description into various modalities of data. -Powered by a universal physics engine re-designed and re-built from the ground up, Genesis integrates various physics solvers and their coupling into a unified framework. This core physics engine is further enhanced by a generative agent framework that operates at an upper level, aiming towards fully **automated data generation** for robotics and beyond. - -Currently, we are open-sourcing the **underlying physics engine and the simulation platform**. Our generative framework is a modular system that incorporates many different generative modules, each handling a certain range of data modalities, routed by a high level agent. Some of the modules integrated existing papers and some are still under submission. Access to our generative feature will be gradually rolled out in the near future. If you are interested, feel free to explore more the [paper list](#papers-behind-genesis) below. - -Genesis is built and will continuously evolve with the following ***long-term missions***: +Genesis aims to: -1. **Lowering the barrier** to using physics simulations and making robotics research accessible to everyone. (See our [commitment](https://genesis-world.readthedocs.io/en/latest/user_guide/overview/mission.html)) -2. **Unifying a wide spectrum of state-of-the-art physics solvers** into a single framework, allowing re-creating the whole physical world in a virtual realm with the highest possible physical, visual and sensory fidelity, using the most advanced simulation techniques. -3. **Minimizing human effort** in collecting and generating data for robotics and other domains, letting the data flywheel spin on its own. +- **Lower the barrier** to using physics simulations, making robotics research accessible to everyone. See our [mission statement](https://genesis-world.readthedocs.io/en/latest/user_guide/overview/mission.html). +- **Unify diverse physics solvers** into a single framework to recreate the physical world with the highest fidelity. +- **Automate data generation**, reducing human effort and letting the data flywheel spin on its own. Project Page: ## Key Features -- **Speed**: Genesis delivers an unprecedented simulation speed -- over 43 million FPS when simulating a Franka robotic arm with a single RTX 4090 (430,000 times faster than real-time). -- **Cross-platform**: Genesis runs natively across different systems (Linux, MacOS, Windows), and across different compute backend (CPU, Nvidia GPU, AMD GPU, Apple Metal). -- **Unification of various physics solvers**: Genesis develops a unified simulation framework that integrates various physics solvers: Rigid body, MPM, SPH, FEM, PBD, Stable Fluid. -- **Support a wide range of material models**: Genesis supports simulation (and the coupling) of rigid and articulated bodies, various types of liquids, gaseous phenomenon, deformable objects, thin-shell objects and granular materials. -- **Support for a wide range of robots**: Robot arm, legged robot, drone, *soft robot*, etc., and extensive support for loading different file types: `MJCF (.xml)`, `URDF`, `.obj`, `.glb`, `.ply`, `.stl`, etc. -- **Photorealistic and high-performance ray-tracer**: Genesis supports native ray-tracing based rendering. -- **Differentiability**: Genesis is designed to be fully compatible with differentiable simulation. Currently, our MPM solver and Tool Solver are differentiable, and differentiability for other solvers will be added soon (starting with rigid-body simulation). -- **Physics-based Tactile Sensor**: Genesis involves a physics-based and differentiable [tactile sensor simulation module](https://github.com/Genesis-Embodied-AI/DiffTactile). This will be integrated to the public version soon (expected in version 0.3.0). -- **User-friendliness**: Genesis is designed in a way to make using simulation as simple as possible. From installation to API design, if there's anything you found counter-intuitive or difficult to use, please [let us know](https://github.com/Genesis-Embodied-AI/Genesis/issues). +- **Speed**: Over 43 million FPS when simulating a Franka robotic arm with a single RTX 4090 (430,000 times faster than real-time). +- **Cross-platform**: Runs on Linux, macOS, Windows, and supports multiple compute backends (CPU, Nvidia/AMD GPUs, Apple Metal). +- **Integration of diverse physics solvers**: Rigid body, MPM, SPH, FEM, PBD, Stable Fluid. +- **Wide range of material models**: Simulation and coupling of rigid bodies, liquids, gases, deformable objects, thin-shell objects, and granular materials. +- **Compatibility with various robots**: Robotic arms, legged robots, drones, *soft robots*, and support for loading `MJCF (.xml)`, `URDF`, `.obj`, `.glb`, `.ply`, `.stl`, and more. +- **Photo-realistic rendering**: Native ray-tracing-based rendering. +- **Differentiability**: Genesis is designed to be fully differentiable. Currently, our MPM solver and Tool Solver support differentiability, with other solvers planned for future versions. +- **Physics-based tactile simulation**: Differentiable [tactile sensor simulation](https://github.com/Genesis-Embodied-AI/DiffTactile) coming soon (expected in version 0.3.0). +- **User-friendliness**: Designed for simplicity, with intuitive installation and APIs. -## Getting Started - -### Quick Installation +## Quick Installation Genesis is available via PyPI: @@ -55,40 +64,63 @@ pip install genesis-world # Requires Python >=3.9; You also need to install **PyTorch** following the [official instructions](https://pytorch.org/get-started/locally/). -If you would like to try out the latest version, we suggest you to git clone from the repo and do `pip install -e .` instead of via PyPI. +For the latest version, clone the repository and install locally: + +```bash +git clone https://github.com/Genesis-Embodied-AI/Genesis.git +cd Genesis +pip install -e . +``` + +## Example Usage -### Documentation +```python +from genesis import Simulation -Please refer to our [documentation site (English)](https://genesis-world.readthedocs.io/en/latest/user_guide/index.html) / [(Chinese)](https://genesis-world.readthedocs.io/zh-cn/latest/user_guide/index.html) for detailed installation steps, tutorials and API references. +# Initialize a simulation +sim = Simulation() + +# Run the simulation +sim.run() +``` + +This simple example demonstrates how to initialize and execute a basic simulation. For more advanced scenarios, refer to our [documentation](https://genesis-world.readthedocs.io/en/latest/user_guide/index.html). + +## Documentation + +Comprehensive documentation is available in [English](https://genesis-world.readthedocs.io/en/latest/user_guide/index.html) and [Chinese](https://genesis-world.readthedocs.io/zh-cn/latest/user_guide/index.html). This includes detailed installation steps, tutorials, and API references. ## Contributing to Genesis -The goal of the Genesis project is to build a fully transparent, user-friendly ecosystem where contributors from both robotics and computer graphics can **come together to collaboratively create a high-efficiency, realistic (both physically and visually) virtual world for robotics research and beyond**. +The Genesis project is an open and collaborative effort. We welcome all forms of contributions from the community, including: -We sincerely welcome *any forms of contributions* from the community to make the world a better place for robots. From **pull requests** for new features, **bug reports**, to even tiny **suggestions** that will make Genesis API more intuitive, all are wholeheartedly appreciated! +- **Pull requests** for new features or bug fixes. +- **Bug reports** through GitHub Issues. +- **Suggestions** to improve Genesis's usability. -## Support +Refer to our [contribution guide](https://github.com/Genesis-Embodied-AI/Genesis/blob/main/CONTRIBUTING.md) for more details. -- Please use Github [Issues](https://github.com/Genesis-Embodied-AI/Genesis/issues) for bug reports and feature requests. +## Support -- Please use GitHub [Discussions](https://github.com/Genesis-Embodied-AI/Genesis/discussions) for discussing ideas, and asking questions. +- Report bugs or request features via GitHub [Issues](https://github.com/Genesis-Embodied-AI/Genesis/issues). +- Join discussions or ask questions on GitHub [Discussions](https://github.com/Genesis-Embodied-AI/Genesis/discussions). -## License and Acknowledgment +## License and Acknowledgments The Genesis source code is licensed under Apache 2.0. -The development of Genesis won't be possible without these amazing open-source projects: -- [Taichi](https://github.com/taichi-dev/taichi): for providing a high-performance cross-platform compute backend. Kudos to all the members providing technical support from taichi! -- [FluidLab](https://github.com/zhouxian/FluidLab) for providing a reference MPM solver implementation -- [SPH_Taichi](https://github.com/erizmr/SPH_Taichi) for providing a reference SPH solver implementation -- [Ten Minute Physics](https://matthias-research.github.io/pages/tenMinutePhysics/index.html) and [PBF3D](https://github.com/WASD4959/PBF3D) for providing a reference PBD solver implementation -- [MuJoCo](https://github.com/google-deepmind/mujoco) and [Brax](https://github.com/google/brax) for providing reference for rigid body dynamics -- [libccd](https://github.com/danfis/libccd) for providing reference for collision detection -- [PyRender](https://github.com/mmatl/pyrender) for rasterization-based renderer -- [LuisaCompute](https://github.com/LuisaGroup/LuisaCompute) and [LuisaRender](https://github.com/LuisaGroup/LuisaRender) for its ray-tracing DSL -- [trimesh](https://github.com/mikedh/trimesh), [PyMeshLab](https://github.com/cnr-isti-vclab/PyMeshLab) and [CoACD](https://github.com/SarahWeiii/CoACD) for geometry processing +Genesis's development has been made possible thanks to these open-source projects: -## Papers behind Genesis +- [Taichi](https://github.com/taichi-dev/taichi): High-performance cross-platform compute backend. Kudos to the Taichi team for their technical support! +- [FluidLab](https://github.com/zhouxian/FluidLab): Reference MPM solver implementation. +- [SPH_Taichi](https://github.com/erizmr/SPH_Taichi): Reference SPH solver implementation. +- [Ten Minute Physics](https://matthias-research.github.io/pages/tenMinutePhysics/index.html) and [PBF3D](https://github.com/WASD4959/PBF3D): Reference PBD solver implementations. +- [MuJoCo](https://github.com/google-deepmind/mujoco): Reference for rigid body dynamics. +- [libccd](https://github.com/danfis/libccd): Reference for collision detection. +- [PyRender](https://github.com/mmatl/pyrender): Rasterization-based renderer. +- [LuisaCompute](https://github.com/LuisaGroup/LuisaCompute) and [LuisaRender](https://github.com/LuisaGroup/LuisaRender): Ray-tracing DSL. + +## Associated Papers Genesis is a large scale effort that integrates state-of-the-art technologies of various existing and on-going research work into a single system. Here we include a non-exhaustive list of all the papers that contributed to the Genesis project in one way or another: @@ -116,7 +148,7 @@ Genesis is a large scale effort that integrates state-of-the-art technologies of ## Citation -If you used Genesis in your research, we would appreciate it if you could cite it. We are still working on a technical report, and before it's public, you could consider citing: +If you use Genesis in your research, please consider citing: ```bibtex @software{Genesis, @@ -126,4 +158,3 @@ If you used Genesis in your research, we would appreciate it if you could cite i year = {2024}, url = {https://github.com/Genesis-Embodied-AI/Genesis} } -``` From 07f1d09e05c5c2a551dbf311056369ba06cd4141 Mon Sep 17 00:00:00 2001 From: 00make Date: Wed, 25 Dec 2024 01:46:24 +0800 Subject: [PATCH 13/35] Updated README_CN and CONTRIBUTING, enhanced readability, and adjusted title format [MD001/heading-increment] --- CONTRIBUTING.md | 19 ++++++++++++----- README_CN.md | 56 ++++++++++++++++++++++++++++++------------------- 2 files changed, 49 insertions(+), 26 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8feaf9cd..6ede5d6c 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -7,9 +7,10 @@ Thank you for your interest in contributing to Genesis! We welcome contributions - [Submitting Code Changes](#submitting-code-changes) - [Reviewing and Merging](#reviewing-and-merging) - [Questions and Discussions](#questions-and-discussions) + --- -### Reporting Bugs +## Reporting Bugs - Before reporting a bug, please search through existing issues to check if it has already been reported. @@ -36,36 +37,44 @@ Thank you for your interest in contributing to Genesis! We welcome contributions Add any other context about the problem here. ``` -### Suggesting Features +## Suggesting Features + - If you have a feature idea, please create an issue labeled `enhancement`. - In the created issue, please provide context, expected outcomes, and potential. -### Submitting Code Changes +## Submitting Code Changes - We use the `pre-commit` configuration to automatically clean up code before committing. Install and run `pre-commit` as follows: 1. Install `pre-commit`: + ```bash pip install pre-commit ``` + 2. Install hooks from the configuration file: + ```bash pre-commit install ``` + After this, `pre-commit` will automatically check and clean up code whenever you make a commit. - (Optional) You can run CI tests locally to ensure you pass the online CI checks. + ```python python -m unittest discover tests ``` + - In the title of your Pull Request, please include [BUG FIX], [FEATURE] or [MISC] to indicate the purpose. - In the description, please provide example code or commands for testing. -### Reviewing and Merging +## Reviewing and Merging - PRs require at least one approval before merging. - Automated checks (e.g., CI tests) must pass. - Use `Squash and Merge` for a clean commit history. -### Questions and Discussions +## Questions and Discussions + - Use [Discussions](https://github.com/Genesis-Embodied-AI/Genesis/discussions) for open-ended topics.