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 enable GPU dynamics & Flatcache + support for both CPU & GPU pipeline #853

Draft
wants to merge 45 commits into
base: og-develop
Choose a base branch
from
Draft
Changes from all commits
Commits
Show all changes
45 commits
Select commit Hold shift + click to select a range
b53f444
Initial implementation of cloth view
cgokmen Aug 19, 2024
2ef872f
Merge branch 'og-develop' into cloth-view
cgokmen Aug 19, 2024
aab7219
Minor fixes
hang-yin Aug 21, 2024
660467a
Some more fixes
hang-yin Aug 21, 2024
cc4de91
Merge branch 'og-develop' of https://github.com/StanfordVL/OmniGibson…
hang-yin Aug 21, 2024
506de18
ParticleView
hang-yin Aug 21, 2024
fc4321c
Fix xformprim and franka bugs
wensi-ai Aug 21, 2024
758315d
Some tensor device fixes
hang-yin Aug 22, 2024
6d65210
Merge branch 'torch-follow-up' of https://github.com/StanfordVL/OmniG…
hang-yin Aug 22, 2024
0e85d1d
Fix joint controller
hang-yin Aug 23, 2024
1629b8c
Fix some device problems
hang-yin Aug 26, 2024
cba26b3
Merge branch 'og-develop' of https://github.com/StanfordVL/OmniGibson…
hang-yin Aug 26, 2024
e35d0bc
Tweek controllers to minimize gpu-cpu sync
hang-yin Aug 27, 2024
6eb723a
Minor fix
hang-yin Aug 27, 2024
e686b39
Refactor controller limit to avoid cpu-gpu sync
hang-yin Aug 28, 2024
4777b5c
More work on minimizing cpu-gpu sync
hang-yin Aug 28, 2024
903963b
Small fix
hang-yin Aug 30, 2024
310e172
Small fix
hang-yin Aug 30, 2024
a8ba9c1
cache some properties for speed
hang-yin Sep 3, 2024
2f102e7
Fix folded unfolded demo with cloth view
hang-yin Sep 5, 2024
01b5bc7
Remove most flatcache and gpu dynamics flags
hang-yin Sep 5, 2024
8d92863
More tensor device fixes
hang-yin Sep 6, 2024
760da56
Merge branch 'og-develop' of https://github.com/StanfordVL/OmniGibson…
hang-yin Sep 11, 2024
43ec4dc
profiling import fix
hang-yin Sep 11, 2024
a2b3f5d
Small fix
hang-yin Sep 11, 2024
a3a5190
Allow cloth prim in cpu pipeline
hang-yin Sep 12, 2024
618476e
Test fixes
hang-yin Sep 12, 2024
ebd0419
Fix cloth prim default positions
hang-yin Sep 12, 2024
1ef6a33
Minor controller fixes
hang-yin Sep 12, 2024
017f413
Temporary fix
hang-yin Sep 13, 2024
4619040
Fix particle requirement
hang-yin Sep 13, 2024
0a93204
Fix data collection type mismatch
hang-yin Sep 13, 2024
a92a135
Refactor tests to run both cpu and gpu modes
hang-yin Sep 14, 2024
08987be
Refactor toggle on test to use primitives IK
hang-yin Sep 14, 2024
ccb37b6
Minor fix
hang-yin Sep 14, 2024
ea676b1
Run both cpu and gpu pipeline tests on CI
hang-yin Sep 14, 2024
ee46579
Fix github test workflow
hang-yin Sep 14, 2024
5d3b11b
Merge branch 'og-develop' of https://github.com/StanfordVL/OmniGibson…
hang-yin Sep 16, 2024
92fb71a
Some bug fixes
hang-yin Sep 17, 2024
8472e29
Fix controller test
hang-yin Sep 17, 2024
5489c24
Minor fixes
hang-yin Sep 17, 2024
196143c
Fix gpu pipeline set velocities warning
hang-yin Sep 17, 2024
f69758c
Merge branch 'og-develop' into cloth-view
cgokmen Dec 4, 2024
499a450
Disable all the things
cgokmen Dec 4, 2024
e3e938b
Convert dof_idx to list for proper indexing into tensor
cgokmen Dec 11, 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
14 changes: 8 additions & 6 deletions .github/workflows/tests.yml
Original file line number Diff line number Diff line change
@@ -19,6 +19,9 @@ jobs:
strategy:
fail-fast: false
matrix:
device:
- cpu
- cuda
test_file:
- test_controllers
- test_curobo
@@ -29,8 +32,7 @@ jobs:
- test_object_removal
- test_object_states
- test_primitives
- test_robot_states_flatcache
- test_robot_states_no_flatcache
- test_robot_states
- test_robot_teleoperation
- test_scene_graph
- test_sensors
@@ -59,21 +61,21 @@ jobs:

- name: Run tests
working-directory: omnigibson-src
run: pytest -s tests/${{ matrix.test_file }}.py --junitxml=${{ matrix.test_file }}.xml && cp ${{ matrix.test_file }}.xml ${GITHUB_WORKSPACE}/
run: pytest -s tests/${{ matrix.test_file }}.py -k ${{ matrix.device }} --junitxml=${{ matrix.test_file }}-${{ matrix.device }}.xml && cp ${{ matrix.test_file }}-${{ matrix.device }}.xml ${GITHUB_WORKSPACE}/
continue-on-error: true

- name: Deploy artifact
- name: Deploy artifacts
uses: actions/upload-artifact@v3
with:
name: ${{ github.run_id }}-tests-${{ matrix.test_file }}
path: ${{ matrix.test_file }}.xml
path: ${{ matrix.test_file }}-${{ matrix.device }}.xml

- name: Check for failures, errors, or missing XML
run: |
if [ ! -f ${{ matrix.test_file }}.xml ]; then
echo "Error: XML file not found, probably due to segfault"
exit 1
elif grep -Eq 'failures="[1-9][0-9]*"|errors="[1-9][0-9]*"' ${{ matrix.test_file }}.xml; then
elif grep -Eq 'failures="[1-9][0-9]*"|errors="[1-9][0-9]*"' ${{ matrix.test_file }}-${{ matrix.device }}.xml; then
echo "Error: Test failures or errors found"
exit 1
else
7 changes: 1 addition & 6 deletions docs/miscellaneous/speed_optimization.md
Original file line number Diff line number Diff line change
@@ -15,12 +15,7 @@ A lot of factors could affect the speed of OmniGibson. Here are a few of them:

1. `ENABLE_HQ_RENDERING`: While it is set to False by default, setting it to True will give us better rendering quality as well as other more advanced rendering features (e.g. isosurface for fluids), but with the cost of dragging down performance.

2. `USE_GPU_DYNAMICS`: setting it to True is required for more advanced features like particles and fluids, but it will lower the performance of OmniGibson.

3. `RENDER_VIEWER_CAMERA`: viewer camera refers to the camera that shows the default viewport in OmniGibson, if you don't want to view the entire scene (e.g. during training), you can set this to False andit will boost OmniGibson performance.

4. `ENABLE_FLATCACHE`: setting it to True will boost OmniGibson performance.

2. `RENDER_VIEWER_CAMERA`: viewer camera refers to the camera that shows the default viewport in OmniGibson, if you don't want to view the entire scene (e.g. during training), you can set this to False andit will boost OmniGibson performance.

## Miscellaneous

65 changes: 38 additions & 27 deletions omnigibson/controllers/controller_base.py
Original file line number Diff line number Diff line change
@@ -4,6 +4,7 @@

import torch as th

import omnigibson as og
from omnigibson.utils.python_utils import Recreatable, Registerable, Serializable, assert_valid_key, classproperty

# Global dicts that will contain mappings
@@ -100,19 +101,25 @@ def __init__(
to the @control_limits entry corresponding to self.control_type
"""
# Store arguments
assert "has_limit" in control_limits, "Expected has_limit specified in control_limits, but does not exist."
assert type(dof_idx) == list, f"Expected dof_idx to be a list, got type {type(dof_idx)} instead."
self._dof_idx = dof_idx
# Store the indices in self.dof_idx that have control limits
self._limited_dof_indices = th.tensor(
[i for i, idx in enumerate(self.dof_idx) if control_limits["has_limit"][idx]],
dtype=th.long,
device=og.sim.device,
)
self._control_freq = control_freq
self._control_limits = {}
# Store control limits as a 3 x 2 x n tensor: [control_type][min/max][dof_idx]
self._control_limits = th.zeros((3, 2, len(dof_idx)), device=og.sim.device)
for motor_type in {"position", "velocity", "effort"}:
if motor_type not in control_limits:
continue

self._control_limits[ControlType.get_type(motor_type)] = [
control_limits[motor_type][0],
control_limits[motor_type][1],
]
assert "has_limit" in control_limits, "Expected has_limit specified in control_limits, but does not exist."
self._dof_has_limits = control_limits["has_limit"]
self._dof_idx = dof_idx.int()
self._control_limits[ControlType.get_type(motor_type)][0] = control_limits[motor_type][0][dof_idx].to(
device=og.sim.device
)
self._control_limits[ControlType.get_type(motor_type)][1] = control_limits[motor_type][1][dof_idx].to(
device=og.sim.device
)

# Generate goal information
self._goal_shapes = self._get_goal_shapes()
@@ -133,8 +140,8 @@ def __init__(
)
command_output_limits = (
(
self._control_limits[self.control_type][0][self.dof_idx],
self._control_limits[self.control_type][1][self.dof_idx],
self._control_limits[self.control_type][0],
self._control_limits[self.control_type][1],
)
if type(command_output_limits) == str and command_output_limits == "default"
else command_output_limits
@@ -169,7 +176,11 @@ def _preprocess_command(self, command):
Array[float]: Processed command vector
"""
# Make sure command is a th.tensor
command = th.tensor([command], dtype=th.float32) if type(command) in {int, float} else command
command = (
th.tensor([command], dtype=th.float32, device=og.sim.device)
if type(command) in {int, float}
else command.to(device=og.sim.device)
)
# We only clip and / or scale if self.command_input_limits exists
if self._command_input_limits is not None:
# Clip
@@ -238,7 +249,9 @@ def update_goal(self, command, control_dict):
), f"Commands must be dimension {self.command_dim}, got dim {len(command)} instead."

# Preprocess and run internal command
self._goal = self._update_goal(command=self._preprocess_command(command), control_dict=control_dict)
self._goal = self._update_goal(
command=self._preprocess_command(command.to(device=og.sim.device)), control_dict=control_dict
)

def _update_goal(self, command, control_dict):
"""
@@ -281,15 +294,13 @@ def clip_control(self, control):
Array[float]: Clipped control signal
"""
clipped_control = control.clip(
self._control_limits[self.control_type][0][self.dof_idx],
self._control_limits[self.control_type][1][self.dof_idx],
self._control_limits[self.control_type][0],
self._control_limits[self.control_type][1],
)
idx = (
self._dof_has_limits[self.dof_idx]
if self.control_type == ControlType.POSITION
else [True] * self.control_dim
)
control[idx] = clipped_control[idx]
if self.control_type == ControlType.POSITION:
control[self._limited_dof_indices] = clipped_control[self._limited_dof_indices]
else:
control = clipped_control
return control

def step(self, control_dict):
@@ -418,12 +429,12 @@ def nums2array(nums, dim):
# Check if input is an Iterable, if so, we simply convert the input to th.tensor and return
# Else, input is a single value, so we map to a numpy array of correct size and return
return (
nums
nums.to(device=og.sim.device)
if isinstance(nums, th.Tensor)
else (
th.tensor(nums, dtype=th.float32)
th.tensor(nums, dtype=th.float32, device=og.sim.device)
if isinstance(nums, Iterable)
else th.ones(dim, dtype=th.float32) * nums
else th.ones(dim, dtype=th.float32, device=og.sim.device) * nums
)
)

@@ -511,7 +522,7 @@ def command_dim(self):
def dof_idx(self):
"""
Returns:
Array[int]: DOF indices corresponding to the specific DOFs being controlled by this robot
th.Tensor[int]: DOF indices corresponding to the specific DOFs being controlled by this robot
"""
return self._dof_idx

9 changes: 5 additions & 4 deletions omnigibson/controllers/dd_controller.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import torch as th

import omnigibson as og
from omnigibson.controllers import ControlType, LocomotionController


@@ -50,8 +51,8 @@ def __init__(
@control_limits velocity limits entry
"""
# Store internal variables
self._wheel_radius = wheel_radius
self._wheel_axle_halflength = wheel_axle_length / 2.0
self._wheel_radius = th.tensor(wheel_radius, device=og.sim.device)
self._wheel_axle_halflength = th.tensor(wheel_axle_length / 2.0, device=og.sim.device)

# If we're using default command output limits, map this to maximum linear / angular velocities
if type(command_output_limits) == str and command_output_limits == "default":
@@ -107,11 +108,11 @@ def compute_control(self, goal_dict, control_dict):
right_wheel_joint_vel = (lin_vel + ang_vel * self._wheel_axle_halflength) / self._wheel_radius

# Return desired velocities
return th.tensor([left_wheel_joint_vel, right_wheel_joint_vel])
return th.stack([left_wheel_joint_vel, right_wheel_joint_vel])

def compute_no_op_goal(self, control_dict):
# This is zero-vector, since we want zero linear / angular velocity
return dict(vel=th.zeros(2))
return dict(vel=th.zeros(2, device=og.sim.device))

def _compute_no_op_action(self, control_dict):
return th.zeros(2, dtype=th.float32)
5 changes: 3 additions & 2 deletions omnigibson/controllers/ik_controller.py
Original file line number Diff line number Diff line change
@@ -2,6 +2,7 @@

import torch as th

import omnigibson as og
import omnigibson.utils.transform_utils as T
from omnigibson.controllers import ControlType, ManipulationController
from omnigibson.controllers.joint_controller import JointController
@@ -331,8 +332,8 @@ def compute_control(self, goal_dict, control_dict):

# Clip values to be within the joint limits
target_joint_pos = target_joint_pos.clamp(
min=self._control_limits[ControlType.get_type("position")][0][self.dof_idx],
max=self._control_limits[ControlType.get_type("position")][1][self.dof_idx],
min=self._control_limits[ControlType.get_type("position")][0],
max=self._control_limits[ControlType.get_type("position")][1],
)

# Optionally pass through smoothing filter for better stability
22 changes: 15 additions & 7 deletions omnigibson/controllers/joint_controller.py
Original file line number Diff line number Diff line change
@@ -2,6 +2,7 @@

import torch as th

import omnigibson as og
import omnigibson.utils.transform_utils as T
from omnigibson.controllers import (
ControlType,
@@ -111,9 +112,13 @@ def __init__(
assert pos_kp is None, "Cannot set pos_kp for JointController with motor_type=effort!"
assert pos_damping_ratio is None, "Cannot set pos_damping_ratio for JointController with motor_type=effort!"
assert vel_kp is None, "Cannot set vel_kp for JointController with motor_type=effort!"
self.pos_kp = pos_kp
self.pos_kd = None if pos_kp is None or pos_damping_ratio is None else 2 * math.sqrt(pos_kp) * pos_damping_ratio
self.vel_kp = vel_kp
self.pos_kp = th.tensor(pos_kp, device=og.sim.device) if pos_kp is not None else None
self.pos_kd = (
None
if pos_kp is None or pos_damping_ratio is None
else th.tensor(2 * math.sqrt(pos_kp) * pos_damping_ratio, device=og.sim.device)
)
self.vel_kp = th.tensor(vel_kp, device=og.sim.device) if vel_kp is not None else None
self._use_impedances = use_impedances
self._use_gravity_compensation = use_gravity_compensation
self._use_cc_compensation = use_cc_compensation
@@ -159,7 +164,10 @@ def _update_goal(self, command, control_dict):

# Compute the final rotations in the quaternion space.
_, end_quat = T.pose_transform(
th.zeros(3), T.euler2quat(delta_rots), th.zeros(3), T.euler2quat(start_rots)
th.zeros(3, dtype=th.float32, device=og.sim.device),
T.euler2quat(delta_rots),
th.zeros(3, dtype=th.float32, device=og.sim.device),
T.euler2quat(start_rots),
)
end_rots = T.quat2euler(end_quat)

@@ -172,8 +180,8 @@ def _update_goal(self, command, control_dict):

# Clip the command based on the limits
target = target.clip(
self._control_limits[ControlType.get_type(self._motor_type)][0][self.dof_idx],
self._control_limits[ControlType.get_type(self._motor_type)][1][self.dof_idx],
self._control_limits[ControlType.get_type(self._motor_type)][0],
self._control_limits[ControlType.get_type(self._motor_type)][1],
)

return dict(target=target)
@@ -238,7 +246,7 @@ def compute_no_op_goal(self, control_dict):
target = control_dict[f"joint_{self._motor_type}"][self.dof_idx]
else:
# For velocity / effort, directly set to 0
target = th.zeros(self.control_dim)
target = th.zeros(self.control_dim, device=og.sim.device)

return dict(target=target)

16 changes: 6 additions & 10 deletions omnigibson/controllers/multi_finger_gripper_controller.py
Original file line number Diff line number Diff line change
@@ -166,13 +166,13 @@ def compute_control(self, goal_dict, control_dict):
should_open = target[0] >= 0.0 if not self._inverted else target[0] > 0.0
if should_open:
u = (
self._control_limits[ControlType.get_type(self._motor_type)][1][self.dof_idx]
self._control_limits[ControlType.get_type(self._motor_type)][1]
if self._open_qpos is None
else self._open_qpos
)
else:
u = (
self._control_limits[ControlType.get_type(self._motor_type)][0][self.dof_idx]
self._control_limits[ControlType.get_type(self._motor_type)][0]
if self._closed_qpos is None
else self._closed_qpos
)
@@ -182,12 +182,8 @@ def compute_control(self, goal_dict, control_dict):

# If we're near the joint limits and we're using velocity / torque control, we zero out the action
if self._motor_type in {"velocity", "torque"}:
violate_upper_limit = (
joint_pos > self._control_limits[ControlType.POSITION][1][self.dof_idx] - self._limit_tolerance
)
violate_lower_limit = (
joint_pos < self._control_limits[ControlType.POSITION][0][self.dof_idx] + self._limit_tolerance
)
violate_upper_limit = joint_pos > self._control_limits[ControlType.POSITION][1] - self._limit_tolerance
violate_lower_limit = joint_pos < self._control_limits[ControlType.POSITION][0] + self._limit_tolerance
violation = th.logical_or(violate_upper_limit * (u > 0), violate_lower_limit * (u < 0))
u *= ~violation

@@ -241,8 +237,8 @@ def _update_grasping_state(self, control_dict):
# Otherwise, the last control signal intends to "move" the gripper
else:
finger_vel = control_dict["joint_velocity"][self.dof_idx]
min_pos = self._control_limits[ControlType.POSITION][0][self.dof_idx]
max_pos = self._control_limits[ControlType.POSITION][1][self.dof_idx]
min_pos = self._control_limits[ControlType.POSITION][0]
max_pos = self._control_limits[ControlType.POSITION][1]

# Make sure we don't have any invalid values (i.e.: fingers should be within the limits)
finger_pos = th.clip(finger_pos, min_pos, max_pos)
17 changes: 13 additions & 4 deletions omnigibson/controllers/osc_controller.py
Original file line number Diff line number Diff line change
@@ -2,6 +2,7 @@

import torch as th

import omnigibson as og
import omnigibson.utils.transform_utils as T
from omnigibson.controllers import ControlType, ManipulationController
from omnigibson.utils.control_utils import orientation_error
@@ -144,9 +145,13 @@ def limiter(target_pos: Array[float], target_quat: Array[float], control_dict: D
)

# Store gains
self.kp = nums2array(nums=kp, dim=6, dtype=th.float32) if kp is not None else None
self.kp = nums2array(nums=kp, dim=6, dtype=th.float32).to(og.sim.device) if kp is not None else None
self.damping_ratio = damping_ratio
self.kp_null = nums2array(nums=kp_null, dim=control_dim, dtype=th.float32) if kp_null is not None else None
self.kp_null = (
nums2array(nums=kp_null, dim=control_dim, dtype=th.float32).to(og.sim.device)
if kp_null is not None
else None
)
self.kd_null = 2 * th.sqrt(self.kp_null) if kp_null is not None else None # critically damped
self.kp_limits = th.tensor(kp_limits, dtype=th.float32)
self.damping_ratio_limits = th.tensor(damping_ratio_limits, dtype=th.float32)
@@ -235,7 +240,8 @@ def limiter(target_pos: Array[float], target_quat: Array[float], control_dict: D
self.decouple_pos_ori = decouple_pos_ori
self.workspace_pose_limiter = workspace_pose_limiter
self.task_name = task_name
self.reset_joint_pos = reset_joint_pos[dof_idx]
self.reset_joint_pos = reset_joint_pos[dof_idx].to(og.sim.device)
self._null_space_identity = th.eye(control_dim, dtype=th.float32, device=og.sim.device)

# Other variables that will be filled in at runtime
self._fixed_quat_target = None
@@ -425,6 +431,7 @@ def compute_control(self, goal_dict, control_dict):
decouple_pos_ori=self.decouple_pos_ori,
base_lin_vel=base_lin_vel,
base_ang_vel=base_ang_vel,
null_space_identity=self._null_space_identity,
).flatten()

# Add gravity compensation
@@ -506,6 +513,7 @@ def _compute_osc_torques(
decouple_pos_ori: bool,
base_lin_vel: th.Tensor,
base_ang_vel: th.Tensor,
null_space_identity: th.Tensor,
):
# Compute the inverse
mm_inv = th.linalg.inv(mm)
@@ -554,8 +562,9 @@ def _compute_osc_torques(
# roboticsproceedings.org/rss07/p31.pdf
if rest_qpos is not None:
j_eef_inv = m_eef @ j_eef @ mm_inv
q_vec = q
u_null = kd_null * -qd + kp_null * ((rest_qpos - q + math.pi) % (2 * math.pi) - math.pi)
u_null = mm @ th.unsqueeze(u_null, dim=-1)
u += (th.eye(control_dim, dtype=th.float32) - j_eef.T @ j_eef_inv) @ u_null
u += (null_space_identity - j_eef.T @ j_eef_inv) @ u_null

return u
Loading