Skip to content

Commit

Permalink
fix formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
@zswang666
zswang666 committed Dec 22, 2024
1 parent bded83d commit 2259023
Show file tree
Hide file tree
Showing 5 changed files with 103 additions and 65 deletions.
130 changes: 77 additions & 53 deletions examples/rigid/set_phys_attr.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -60,103 +60,127 @@ 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())

# 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)

Expand Down
8 changes: 6 additions & 2 deletions genesis/engine/entities/rigid_entity/rigid_entity.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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):
Expand Down
10 changes: 6 additions & 4 deletions genesis/engine/solvers/rigid/constraint_solver_decomp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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):
Expand Down
4 changes: 3 additions & 1 deletion genesis/engine/solvers/rigid/constraint_solver_decomp_island.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
16 changes: 11 additions & 5 deletions genesis/engine/solvers/rigid/rigid_solver_decomp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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
)
Expand Down Expand Up @@ -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":
Expand Down Expand Up @@ -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(())
Expand Down

0 comments on commit 2259023

Please sign in to comment.