Skip to content

Commit

Permalink
bug fix with setting before hide poses
Browse files Browse the repository at this point in the history
  • Loading branch information
@StoneT2000
StoneT2000 committed Mar 4, 2024
1 parent 33a0256 commit 0740017
Show file tree
Hide file tree
Showing 4 changed files with 100 additions and 90 deletions.
8 changes: 4 additions & 4 deletions mani_skill2/utils/structs/actor.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

from dataclasses import dataclass
from functools import cache
from typing import TYPE_CHECKING, List, Literal, Union

import numpy as np
Expand Down Expand Up @@ -140,6 +141,7 @@ def set_state(self, state: Array):
self.set_linear_velocity(state[7:10])
self.set_angular_velocity(state[10:13])

@cache
def has_collision_shapes(self):
return (
len(
Expand All @@ -165,6 +167,7 @@ def hide_visual(self):
self.before_hide_pose = self.px.cuda_rigid_body_data.torch()[
self._body_data_index, :7
].clone()

temp_pose = self.pose.raw_pose
temp_pose[..., :3] += 99999
self.pose = temp_pose
Expand Down Expand Up @@ -228,9 +231,6 @@ def pose(self) -> Pose:
raw_pose = self.px.cuda_rigid_body_data.torch()[
self._body_data_index, :7
]
# if self.hidden:
# print(self.name, "hidden", raw_pose[0, :3])
# raw_pose[..., :3] = raw_pose[..., :3] - 99999
return Pose.create(raw_pose)
else:
assert len(self._objs) == 1
Expand All @@ -242,7 +242,7 @@ def pose(self, arg1: Union[Pose, sapien.Pose, Array]) -> None:
if not isinstance(arg1, torch.Tensor):
arg1 = vectorize_pose(arg1)
if self.hidden:
self.before_hide_pose = arg1
self.before_hide_pose[self._scene._reset_mask] = arg1
else:
self.px.cuda_rigid_body_data.torch()[
self._body_data_index[self._scene._reset_mask], :7
Expand Down
Binary file renamed ...45b5-cp310-cp310-manylinux2014_x86_64.whl → ...1b1f-cp310-cp310-manylinux2014_x86_64.whl
View file
Binary file not shown.
Binary file renamed ...a245b5-cp39-cp39-manylinux2014_x86_64.whl → ...781b1f-cp39-cp39-manylinux2014_x86_64.whl
View file
Binary file not shown.
182 changes: 96 additions & 86 deletions tests/test_gpu_envs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -269,95 +269,105 @@ def test_hidden_objs(env_id):
env_id, num_envs=16, vectorization_mode="custom"
)
obs, _ = env.reset()

# for PickCube, this is env.goal_site
hide_obj = env.unwrapped._hidden_objects[0]

raw_pose = hide_obj.pose.raw_pose.clone()
p = hide_obj.pose.p.clone()
q = hide_obj.pose.q.clone()
linvel = hide_obj.linear_velocity.clone()
angvel = hide_obj.angular_velocity.clone()

# hide_visual tests
def test_fn():
# note that we use torch.is_close instead of checking exact match as data can change a tiny amount (1e-7 range ish)
# when applying then fetching new poses

# for PickCube, this is env.goal_site
raw_pose = hide_obj.pose.raw_pose.clone()
p = hide_obj.pose.p.clone()
q = hide_obj.pose.q.clone()
linvel = hide_obj.linear_velocity.clone()
angvel = hide_obj.angular_velocity.clone()
# 1. check relevant hidden properties are active
assert hide_obj.hidden
assert hasattr(hide_obj, "before_hide_pose")

# 2. check state data for new pos is not too low or high
assert (
hide_obj.px.cuda_rigid_body_data.torch()[
hide_obj._body_data_index, :7
].clone()[..., :3]
> 1e3
).all()
assert (
hide_obj.px.cuda_rigid_body_data.torch()[
hide_obj._body_data_index, :7
].clone()[..., :3]
< 1e6
).all()

# 3. check that linvel and angvel same as before
assert (hide_obj.linear_velocity == linvel).all()
assert (hide_obj.angular_velocity == angvel).all()

# 4. Check data stored in buffer has same q but different p
assert (
hide_obj.px.cuda_rigid_body_data.torch()[
hide_obj._body_data_index, :7
].clone()[..., :3]
!= p
).all()
assert (
hide_obj.px.cuda_rigid_body_data.torch()[
hide_obj._body_data_index, :7
].clone()[..., 3:]
== q
).all()

# 5. Check data stored in before_hide_pose has same q and p
assert (hide_obj.before_hide_pose[..., :3] == p).all()
assert (hide_obj.before_hide_pose[..., 3:] == q).all()

# 6. check that direct calls to raw_pose, pos, and rot same as before
# (should return `before_hide_pose`)
assert (hide_obj.pose.raw_pose == raw_pose).all()
assert (hide_obj.pose.p == p).all()
assert (hide_obj.pose.q == q).all()
assert (hide_obj.pose.raw_pose == hide_obj.before_hide_pose).all()

# show_visual tests
hide_obj.show_visual()

# 1. check relevant hidden properties are active
assert not hide_obj.hidden

# 2. check that qvel, linvel, angvel same as before
assert (hide_obj.linear_velocity == linvel).all()
assert (hide_obj.angular_velocity == angvel).all()

# 3. check gpu buffer goes back to normal
assert torch.isclose(
hide_obj.px.cuda_rigid_body_data.torch()[
hide_obj._body_data_index, :7
].clone()[..., :3],
p,
).all()
assert torch.isclose(
hide_obj.px.cuda_rigid_body_data.torch()[
hide_obj._body_data_index, :7
].clone()[..., 3:],
q,
).all()

# 4. check that direct calls to raw_pose, pos, and rot same as before
assert torch.isclose(hide_obj.pose.raw_pose, raw_pose).all()
assert torch.isclose(hide_obj.pose.p, p).all()
assert torch.isclose(hide_obj.pose.q, q).all()

# Test after reset
hide_obj.hide_visual()

# 1. check relevant hidden properties are active
assert hide_obj.hidden
assert hasattr(hide_obj, "before_hide_pose")

# 2. check state data for new pos is not too low or high
assert (
hide_obj.px.cuda_rigid_body_data.torch()[hide_obj._body_data_index, :7].clone()[
..., :3
]
> 1e3
).all()
assert (
hide_obj.px.cuda_rigid_body_data.torch()[hide_obj._body_data_index, :7].clone()[
..., :3
]
< 1e6
).all()

# 3. check that linvel and angvel same as before
assert (hide_obj.linear_velocity == linvel).all()
assert (hide_obj.angular_velocity == angvel).all()

# 4. Check data stored in buffer has same q but different p
assert (
hide_obj.px.cuda_rigid_body_data.torch()[hide_obj._body_data_index, :7].clone()[
..., :3
]
!= p
).all()
assert (
hide_obj.px.cuda_rigid_body_data.torch()[hide_obj._body_data_index, :7].clone()[
..., 3:
]
== q
).all()

# 5. Check data stored in before_hide_pose has same q and p
assert (hide_obj.before_hide_pose[..., :3] == p).all()
assert (hide_obj.before_hide_pose[..., 3:] == q).all()

# 6. check that direct calls to raw_pose, pos, and rot same as before
# (should return `before_hide_pose`)
assert (hide_obj.pose.raw_pose == raw_pose).all()
assert (hide_obj.pose.p == p).all()
assert (hide_obj.pose.q == q).all()
assert (hide_obj.pose.raw_pose == hide_obj.before_hide_pose).all()

# show_visual tests
hide_obj.show_visual()

# 1. check relevant hidden properties are active
assert not hide_obj.hidden

# 2. check that qvel, linvel, angvel same as before
assert (hide_obj.linear_velocity == linvel).all()
assert (hide_obj.angular_velocity == angvel).all()

# 3. check gpu buffer goes back to normal
assert (
hide_obj.px.cuda_rigid_body_data.torch()[hide_obj._body_data_index, :7].clone()[
..., :3
]
== p
).all()
assert (
hide_obj.px.cuda_rigid_body_data.torch()[hide_obj._body_data_index, :7].clone()[
..., 3:
]
== q
).all()

# 4. check that direct calls to raw_pose, pos, and rot same as before
assert (hide_obj.pose.raw_pose == raw_pose).all()
assert (hide_obj.pose.p == p).all()
assert (hide_obj.pose.q == q).all()

test_fn()
# Test after partial resets
hide_obj.hide_visual()
env_idx = torch.arange(16, dtype=int, device=env.device)
reset_mask = torch.zeros(16, dtype=bool, device=env.device)
for i in [1, 3, 7, 11]:
reset_mask[i] = True
env.reset(options=dict(env_idx=env_idx[reset_mask]))
test_fn()
env.close()
del env

Expand Down

0 comments on commit 0740017

Please sign in to comment.