Skip to content
New issue

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

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

Already on GitHub? Sign in to your account

[WIP] Always use GPU pipeline and flatcache #319

Closed
wants to merge 31 commits into from
Closed
Changes from 1 commit
Commits
Show all changes
31 commits
Select commit Hold shift + click to select a range
d371373
Remove ENABLE_FLATCACHE flag
cgokmen Nov 9, 2023
cdd00d3
Implement new cloth particle getter
cgokmen Nov 9, 2023
82aac5d
Always enable flatcache
cgokmen Nov 9, 2023
cdedf29
Make flatcache non-optional
cgokmen Nov 9, 2023
b19ca49
Remove unused stuff
cgokmen Nov 9, 2023
49c2ddf
Fix bad import
cgokmen Nov 9, 2023
be50d52
Fix issue with premature render breaking things
cgokmen Nov 10, 2023
b954fdd
Default GPU dynamics to True for now
cgokmen Nov 11, 2023
722ec9d
Prepare cloth prim with new GPU approach
cgokmen Nov 11, 2023
25f6727
Use GPU pipeline, dodge bug in stage callback
cgokmen Nov 11, 2023
c40e949
Fix some bugs in cloth remeshing pipeline
cgokmen Nov 11, 2023
b1ef448
More fixes to remeshing system
cgokmen Nov 11, 2023
e8b3d65
Some debugging code
cgokmen Nov 11, 2023
641111a
Update cloth_prim.py
cgokmen Nov 13, 2023
2338582
No more enable flatcache flag
cgokmen Nov 13, 2023
2d621e0
Merge branch 'og-develop' into core-classes
cgokmen Dec 22, 2023
d44e554
Make all the tensors into numpy arrays
cgokmen Dec 22, 2023
c208bc4
Revert "Make all the tensors into numpy arrays"
cgokmen Dec 22, 2023
c087a3c
Get GPU pipeline working again
cgokmen Dec 23, 2023
51838b5
Merge branch 'og-develop' into core-classes
cgokmen Feb 23, 2024
ee3e2f8
Remove GPU dynamics flag, unnecessary flag changes, FlatcacheAPI
cgokmen Feb 23, 2024
a4b5bf9
Fix bad lazy omni calls
cgokmen Feb 23, 2024
a58d70c
Remove eager omni import
cgokmen Feb 23, 2024
965707b
More lazy imports
cgokmen Feb 23, 2024
3a24e3d
Fix merge error
cgokmen Feb 23, 2024
6cd67d1
More fixes
cgokmen Feb 23, 2024
83f3e30
Use GPU-compatible interface for velocity
cgokmen Feb 23, 2024
47f72b0
No kinematic objects for now
cgokmen Feb 24, 2024
a213c6d
Temporary mods to scene selector
cgokmen Feb 24, 2024
767a946
Revert "No more enable flatcache flag"
cgokmen Feb 29, 2024
77f676c
Merge branch 'core-classes' of https://github.com/StanfordVL/OmniGibs…
cgokmen Feb 29, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Revert "Make all the tensors into numpy arrays"
This reverts commit d44e554.
cgokmen committed Dec 22, 2023

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
commit c208bc4bc0955de91fae55b35e8abd58f6ee552a
2 changes: 1 addition & 1 deletion omnigibson/prims/cloth_prim.py
Original file line number Diff line number Diff line change
@@ -183,7 +183,7 @@ def set_particle_positions(self, positions, idxs=None):
f"Got mismatch in particle setting size: {len(positions)}, vs. number of expected particles {n_expected}!"

# First, get the particle positions.
cur_pos = self._cloth_prim_view.get_world_positions().cpu().numpy()
cur_pos = self._cloth_prim_view.get_world_positions()

# Then apply the new positions at the appropriate indices
cur_pos[0, idxs] = torch.Tensor(positions, device=cur_pos.device)
30 changes: 11 additions & 19 deletions omnigibson/prims/entity_prim.py
Original file line number Diff line number Diff line change
@@ -796,7 +796,7 @@ def get_joint_positions(self, normalized=False):
# Run sanity checks -- make sure we are articulated
assert self.n_joints > 0, "Tried to call method not intended for entity prim with no joints!"

joint_positions = self._articulation_view.get_joint_positions().cpu().numpy().reshape(self.n_dof)
joint_positions = self._articulation_view.get_joint_positions().reshape(self.n_dof)

# Possibly normalize values when returning
return self._normalize_positions(positions=joint_positions) if normalized else joint_positions
@@ -814,7 +814,7 @@ def get_joint_velocities(self, normalized=False):
# Run sanity checks -- make sure we are articulated
assert self.n_joints > 0, "Tried to call method not intended for entity prim with no joints!"

joint_velocities = self._articulation_view.get_joint_velocities().cpu().numpy().reshape(self.n_dof)
joint_velocities = self._articulation_view.get_joint_velocities().reshape(self.n_dof)

# Possibly normalize values when returning
return self._normalize_velocities(velocities=joint_velocities) if normalized else joint_velocities
@@ -832,7 +832,7 @@ def get_joint_efforts(self, normalized=False):
# Run sanity checks -- make sure we are articulated
assert self.n_joints > 0, "Tried to call method not intended for entity prim with no joints!"

joint_efforts = self._articulation_view.get_applied_joint_efforts().cpu().numpy().reshape(self.n_dof)
joint_efforts = self._articulation_view.get_applied_joint_efforts().reshape(self.n_dof)

# Possibly normalize values when returning
return self._normalize_efforts(efforts=joint_efforts) if normalized else joint_efforts
@@ -889,12 +889,8 @@ def get_position_orientation(self):
if self._articulation_view is None:
return self.root_link.get_position_orientation()

pos, ori = self._articulation_view.get_world_poses()
pos = pos.cpu().numpy()
ori = ori.cpu().numpy()
assert np.isclose(np.linalg.norm(ori), 1, atol=1e-3), \
f"{self.prim_path} orientation {ori} is not a unit quaternion."
return pos[0], ori[0][[1, 2, 3, 0]]
positions, orientations = self._articulation_view.get_world_poses()
return positions[0], orientations[0][[1, 2, 3, 0]]

def set_local_pose(self, position=None, orientation=None):
# Delegate to RigidPrim if we are not articulated
@@ -913,12 +909,8 @@ def get_local_pose(self):
if self._articulation_view is None:
return self.root_link.get_local_pose()

pos, ori = self._articulation_view.get_local_poses()
pos = pos.cpu().numpy()
ori = ori.cpu().numpy()
assert np.isclose(np.linalg.norm(ori), 1, atol=1e-3), \
f"{self.prim_path} orientation {ori} is not a unit quaternion."
return pos[0], ori[0][[1, 2, 3, 0]]
positions, orientations = self._articulation_view.get_local_poses()
return positions[0], orientations[0][[1, 2, 3, 0]]

# TODO: Is the omni joint damping (used for driving motors) same as dissipative joint damping (what we had in pb)?
@property
@@ -1192,23 +1184,23 @@ def get_coriolis_and_centrifugal_forces(self):
n-array: (N,) shaped per-DOF coriolis and centrifugal forces experienced by the entity, if articulated
"""
assert self.articulated, "Cannot get coriolis and centrifugal forces for non-articulated entity!"
return self._articulation_view.get_coriolis_and_centrifugal_forces().cpu().numpy().reshape(self.n_dof)
return self._articulation_view.get_coriolis_and_centrifugal_forces().reshape(self.n_dof)

def get_generalized_gravity_forces(self):
"""
Returns:
n-array: (N, N) shaped per-DOF gravity forces, if articulated
"""
assert self.articulated, "Cannot get generalized gravity forces for non-articulated entity!"
return self._articulation_view.get_generalized_gravity_forces().cpu().numpy().reshape(self.n_dof)
return self._articulation_view.get_generalized_gravity_forces().reshape(self.n_dof)

def get_mass_matrix(self):
"""
Returns:
n-array: (N, N) shaped per-DOF mass matrix, if articulated
"""
assert self.articulated, "Cannot get mass matrix for non-articulated entity!"
return self._articulation_view.get_mass_matrices().cpu().numpy().reshape(self.n_dof, self.n_dof)
return self._articulation_view.get_mass_matrices().reshape(self.n_dof, self.n_dof)

def get_jacobian(self):
"""
@@ -1218,7 +1210,7 @@ def get_jacobian(self):
(i.e.: there is an additional "floating" joint tying the robot to the world frame)
"""
assert self.articulated, "Cannot get jacobian for non-articulated entity!"
return self._articulation_view.get_jacobians().cpu().numpy().squeeze(axis=0)
return self._articulation_view.get_jacobians().squeeze(axis=0)

def get_relative_jacobian(self):
"""
26 changes: 13 additions & 13 deletions omnigibson/prims/joint_prim.py
Original file line number Diff line number Diff line change
@@ -148,7 +148,7 @@ def _initialize(self):
if self.articulated:
control_types = []
stiffnesses, dampings = self._articulation_view.get_gains(joint_indices=self.dof_indices)
for i, (kp, kd) in enumerate(zip(stiffnesses[0].cpu().numpy(), dampings[0].cpu().numpy())):
for i, (kp, kd) in enumerate(zip(stiffnesses[0], dampings[0])):
# Infer control type based on whether kp and kd are 0 or not, as well as whether this joint is driven or not
# TODO: Maybe assert mutual exclusiveness here?
if not self._driven:
@@ -335,7 +335,7 @@ def max_velocity(self):
# Only support revolute and prismatic joints for now
assert self.is_single_dof, "Joint properties only supported for a single DOF currently!"
# We either return the raw value or a default value if there is no max specified
raw_vel = self._articulation_view.get_max_velocities(joint_indices=self.dof_indices).cpu().numpy()[0][0]
raw_vel = self._articulation_view.get_max_velocities(joint_indices=self.dof_indices)[0][0]
default_max_vel = m.DEFAULT_MAX_REVOLUTE_VEL if self.joint_type == JointType.JOINT_REVOLUTE else m.DEFAULT_MAX_PRISMATIC_VEL
return default_max_vel if raw_vel is None or np.abs(raw_vel) > m.INF_VEL_THRESHOLD else raw_vel

@@ -362,7 +362,7 @@ def max_effort(self):
# Only support revolute and prismatic joints for now
assert self.is_single_dof, "Joint properties only supported for a single DOF currently!"
# We either return the raw value or a default value if there is no max specified
raw_effort = self._articulation_view.get_max_efforts(joint_indices=self.dof_indices).cpu().numpy()[0][0]
raw_effort = self._articulation_view.get_max_efforts(joint_indices=self.dof_indices)[0][0]
return m.DEFAULT_MAX_EFFORT if raw_effort is None or np.abs(raw_effort) > m.INF_EFFORT_THRESHOLD else raw_effort

@max_effort.setter
@@ -388,7 +388,7 @@ def stiffness(self):
# Only support revolute and prismatic joints for now
assert self.is_single_dof, "Joint properties only supported for a single DOF currently!"
stiffnesses, _ = self._articulation_view.get_gains(joint_indices=self.dof_indices)[0]
return stiffnesses.cpu().numpy()[0][0]
return stiffnesses[0][0]

@stiffness.setter
def stiffness(self, stiffness):
@@ -413,7 +413,7 @@ def damping(self):
# Only support revolute and prismatic joints for now
assert self.is_single_dof, "Joint properties only supported for a single DOF currently!"
_, dampings = self._articulation_view.get_gains(joint_indices=self.dof_indices)[0]
return dampings.cpu().numpy()[0][0]
return dampings[0][0]

@damping.setter
def damping(self, damping):
@@ -435,7 +435,7 @@ def friction(self):
Returns:
float: friction for this joint
"""
return self._articulation_view.get_friction_coefficients(joint_indices=self.dof_indices).cpu().numpy()[0][0]
return self._articulation_view.get_friction_coefficients(joint_indices=self.dof_indices)[0][0]

@friction.setter
def friction(self, friction):
@@ -459,7 +459,7 @@ def lower_limit(self):
# Only support revolute and prismatic joints for now
assert self.is_single_dof, "Joint properties only supported for a single DOF currently!"
# We either return the raw value or a default value if there is no max specified
raw_pos_lower, raw_pos_upper = self._articulation_view.get_joint_limits(joint_indices=self.dof_indices).cpu().numpy().flatten()
raw_pos_lower, raw_pos_upper = self._articulation_view.get_joint_limits(joint_indices=self.dof_indices).flatten()
return -m.DEFAULT_MAX_POS \
if raw_pos_lower is None or raw_pos_lower == raw_pos_upper or np.abs(raw_pos_lower) > m.INF_POS_THRESHOLD \
else raw_pos_lower
@@ -487,7 +487,7 @@ def upper_limit(self):
# Only support revolute and prismatic joints for now
assert self.is_single_dof, "Joint properties only supported for a single DOF currently!"
# We either return the raw value or a default value if there is no max specified
raw_pos_lower, raw_pos_upper = self._articulation_view.get_joint_limits(joint_indices=self.dof_indices).cpu().numpy().flatten()
raw_pos_lower, raw_pos_upper = self._articulation_view.get_joint_limits(joint_indices=self.dof_indices).flatten()
return m.DEFAULT_MAX_POS \
if raw_pos_upper is None or raw_pos_lower == raw_pos_upper or np.abs(raw_pos_upper) > m.INF_POS_THRESHOLD \
else raw_pos_upper
@@ -512,7 +512,7 @@ def has_limit(self):
"""
# Only support revolute and prismatic joints for now
assert self.is_single_dof, "Joint properties only supported for a single DOF currently!"
return np.all(np.abs(self._articulation_view.get_joint_limits(joint_indices=self.dof_indices).cpu().numpy()) < m.INF_POS_THRESHOLD)
return np.all(np.abs(self._articulation_view.get_joint_limits(joint_indices=self.dof_indices)) < m.INF_POS_THRESHOLD)

@property
def axis(self):
@@ -598,9 +598,9 @@ def get_state(self, normalized=False):
assert self.articulated, "Can only get state for articulated joints!"

# Grab raw states
pos = self._articulation_view.get_joint_positions(joint_indices=self.dof_indices).cpu().numpy()[0]
vel = self._articulation_view.get_joint_velocities(joint_indices=self.dof_indices).cpu().numpy()[0]
effort = self._articulation_view.get_applied_joint_efforts(joint_indices=self.dof_indices).cpu().numpy()[0]
pos = self._articulation_view.get_joint_positions(joint_indices=self.dof_indices)[0]
vel = self._articulation_view.get_joint_velocities(joint_indices=self.dof_indices)[0]
effort = self._articulation_view.get_applied_joint_efforts(joint_indices=self.dof_indices)[0]

# Potentially normalize if requested
if normalized:
@@ -624,7 +624,7 @@ def get_target(self, normalized=False):
assert self.articulated, "Can only get targets for articulated joints!"

# Grab raw states
targets = self._articulation_view.get_applied_actions().cpu().numpy()
targets = self._articulation_view.get_applied_actions()
pos = targets.joint_positions[0][self.dof_indices]
vel = targets.joint_velocities[0][self.dof_indices]

21 changes: 8 additions & 13 deletions omnigibson/prims/rigid_prim.py
Original file line number Diff line number Diff line change
@@ -249,7 +249,7 @@ def get_linear_velocity(self):
Returns:
np.ndarray: current linear velocity of the the rigid prim. Shape (3,).
"""
return self._rigid_prim_view.get_linear_velocities().cpu().numpy()[0]
return self._rigid_prim_view.get_linear_velocities()[0]

def set_angular_velocity(self, velocity):
"""
@@ -265,7 +265,7 @@ def get_angular_velocity(self):
Returns:
np.ndarray: current angular velocity of the the rigid prim. Shape (3,).
"""
return self._rigid_prim_view.get_angular_velocities().cpu().numpy()[0]
return self._rigid_prim_view.get_angular_velocities()[0]

def set_position_orientation(self, position=None, orientation=None):
if position is not None:
@@ -279,8 +279,7 @@ def set_position_orientation(self, position=None, orientation=None):

def get_position_orientation(self):
pos, ori = self._rigid_prim_view.get_world_poses()
pos = pos.cpu().numpy()
ori = ori.cpu().numpy()

assert np.isclose(np.linalg.norm(ori), 1, atol=1e-3), \
f"{self.prim_path} orientation {ori} is not a unit quaternion."
return pos[0], ori[0][[1, 2, 3, 0]]
@@ -294,12 +293,8 @@ def set_local_pose(self, position=None, orientation=None):
BoundingBoxAPI.clear()

def get_local_pose(self):
pos, ori = self._rigid_prim_view.get_local_poses()
pos = pos.cpu().numpy()
ori = ori.cpu().numpy()
assert np.isclose(np.linalg.norm(ori), 1, atol=1e-3), \
f"{self.prim_path} orientation {ori} is not a unit quaternion."
return pos[0], ori[0][[1, 2, 3, 0]]
positions, orientations = self._rigid_prim_view.get_local_poses()
return positions[0], orientations[0][[1, 2, 3, 0]]

@property
def _rigid_prim_view(self):
@@ -404,7 +399,7 @@ def mass(self):
Returns:
float: mass of the rigid body in kg.
"""
mass = self._rigid_prim_view.get_masses().cpu().numpy()[0]
mass = self._rigid_prim_view.get_masses()[0]

# Fallback to analytical computation of volume * density
if mass == 0:
@@ -426,15 +421,15 @@ def density(self):
Returns:
float: density of the rigid body in kg / m^3.
"""
raw_usd_mass = self._rigid_prim_view.get_masses().cpu().numpy()[0]
raw_usd_mass = self._rigid_prim_view.get_masses()[0]
# We first check if the raw usd mass is specified, since mass overrides density
# If it's specified, we infer density based on that value divided by volume
# Otherwise, we try to directly grab the raw usd density value, and if that value
# does not exist, we return 1000 since that is the canonical density assigned by omniverse
if raw_usd_mass != 0:
density = raw_usd_mass / self.volume
else:
density = self._rigid_prim_view.get_densities().cpu().numpy()[0]
density = self._rigid_prim_view.get_densities()[0]
if density == 0:
density = 1000.0

4 changes: 2 additions & 2 deletions omnigibson/systems/macro_particle_system.py
Original file line number Diff line number Diff line change
@@ -1179,7 +1179,7 @@ def add_particle(cls, prim_path, scale, idn=None):
def get_particles_position_orientation(cls):
# Note: This gets the center of the sphere approximation of the particles, NOT the actual particle frames!
if cls.n_particles > 0:
tfs = cls.particles_view.get_transforms().cpu().numpy()
tfs = cls.particles_view.get_transforms()
pos, ori = tfs[:, :3], tfs[:, 3:]
pos = pos + T.quat2mat(ori) @ cls._particle_offset
else:
@@ -1242,7 +1242,7 @@ def get_particles_velocities(cls):
- (n, 3)-array: per-particle (ax, ay, az) angular velocities in the world frame
"""
if cls.n_particles > 0:
vels = cls.particles_view.get_velocities().cpu().numpy()
vels = cls.particles_view.get_velocities()
lin_vel, ang_vel = vels[:, :3], vels[:, 3:]
else:
lin_vel, ang_vel = np.array([]).reshape(0, 3), np.array([]).reshape(0, 3)