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

Optional Numpy Compute Backend #1051

Merged
merged 49 commits into from
Dec 19, 2024
Merged
Changes from 2 commits
Commits
Show all changes
49 commits
Select commit Hold shift + click to select a range
06d1578
refactor backend to use np for better single-threaded efficiency
cremebrule Dec 3, 2024
38b2a47
fix small typo
cremebrule Dec 3, 2024
429cfeb
update setup requirements
cremebrule Dec 3, 2024
8fa44a0
cache indexing values for better efficiency
cremebrule Dec 4, 2024
5d3bce5
improve control loop efficiency, cache relative tf computations
cremebrule Dec 4, 2024
6787df3
optimize controller computation functions
cremebrule Dec 5, 2024
692df78
cache properties that should not change during runtime
cremebrule Dec 5, 2024
d092a6f
optimize np numba transform functions
cremebrule Dec 5, 2024
4ef6f4d
Merge branch 'asset-conversion' into feat/np-opt
cremebrule Dec 12, 2024
69df439
refactor controller backend and general compute backend
cremebrule Dec 12, 2024
3b8d114
update filters to use compute backend
cremebrule Dec 12, 2024
e45e719
fix controller bugs
cremebrule Dec 12, 2024
e9219c2
fix proprioception dict generation
cremebrule Dec 12, 2024
50d390a
fix holonomic robot pose setting
cremebrule Dec 12, 2024
fd07d8a
fix controllable object and filters state de/serialization
cremebrule Dec 12, 2024
9401922
make np tf use numba
cremebrule Dec 12, 2024
920d161
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Dec 12, 2024
5121aa5
ensure actions and control dict is handled properly externally
cremebrule Dec 12, 2024
60d65ab
Merge remote-tracking branch 'origin/feat/np-opt' into feat/np-opt
cremebrule Dec 12, 2024
8d6f25f
fix control torch reference
cremebrule Dec 12, 2024
1dd3b28
Merge branch 'asset-conversion' into feat/np-opt
cremebrule Dec 12, 2024
d4553eb
Merge branch 'asset-conversion' into feat/np-opt
cremebrule Dec 12, 2024
f73f9f0
Merge branch 'asset-conversion' into feat/np-opt
cremebrule Dec 12, 2024
64fab4f
Merge branch 'asset-conversion' into feat/np-opt
cremebrule Dec 12, 2024
62d08b9
use compute backend for filter product
cremebrule Dec 12, 2024
5d672a6
use 0-dim command for null joint controller
cremebrule Dec 12, 2024
0fca8d6
improve curobo test determinism
cremebrule Dec 12, 2024
c83772d
make collision activation settable for better reproducability in curo…
cremebrule Dec 12, 2024
11b9876
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Dec 12, 2024
8268247
Merge branch 'asset-conversion' into feat/np-opt
cremebrule Dec 12, 2024
a546d45
Merge remote-tracking branch 'origin/feat/np-opt' into feat/np-opt
cremebrule Dec 12, 2024
b606001
fix IK controller arg parsing
cremebrule Dec 12, 2024
474aa0d
update comments external -> custom dataset
cremebrule Dec 12, 2024
ee78e94
Merge branch 'asset-conversion' into feat/np-opt
cremebrule Dec 12, 2024
04c37c5
update version and remove extraneous np import
cremebrule Dec 12, 2024
3e015f7
Merge branch 'og-develop' into feat/np-opt
cremebrule Dec 12, 2024
e31c12a
make urdf paths relative during custom asset import
cremebrule Dec 13, 2024
6275ab9
handle vision sensor cleanup properly
cremebrule Dec 13, 2024
42e1d9d
fix default command output behavior for joint controllers
cremebrule Dec 14, 2024
2f856aa
refactor compute backend to allow for explicit externally-defined fun…
cremebrule Dec 19, 2024
a2d35c9
fix default command output generation for multi finger controller
cremebrule Dec 19, 2024
475cbc7
fix semantics for null joint controller
cremebrule Dec 19, 2024
8faa0a9
verify that th and np transform utils have 1-to-1 corresponding funct…
cremebrule Dec 19, 2024
1c67bba
update null joint controller default args in robot classes
cremebrule Dec 19, 2024
66cfff7
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Dec 19, 2024
374513b
re-add quaternions_close
cremebrule Dec 19, 2024
209a8b3
Merge remote-tracking branch 'origin/feat/np-opt' into feat/np-opt
cremebrule Dec 19, 2024
46c6ecb
re-add transformation_matrix
cremebrule Dec 19, 2024
a1fbadd
add copysign to np transform utils
cremebrule Dec 19, 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
8 changes: 3 additions & 5 deletions omnigibson/controllers/controller_base.py
Original file line number Diff line number Diff line change
@@ -3,8 +3,9 @@
from enum import IntEnum

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

from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.utils.python_utils import Recreatable, Registerable, Serializable, assert_valid_key, classproperty

# Global dicts that will contain mappings
REGISTERED_CONTROLLERS = dict()
@@ -34,9 +35,6 @@ class IsGraspingState(IntEnum):
FALSE = -1





# Define macros
class ControlType:
NONE = -1
@@ -453,7 +451,7 @@ def nums2array(nums, dim):
return (
nums
if isinstance(nums, cb.arr_type)
else cb.array(nums)if isinstance(nums, Iterable) else cb.ones(dim) * nums
else cb.array(nums) if isinstance(nums, Iterable) else cb.ones(dim) * nums
)

@property
17 changes: 13 additions & 4 deletions omnigibson/controllers/ik_controller.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
import math
from collections.abc import Iterable
from omnigibson.utils.backend_utils import _ComputeBackend, _ComputeTorchBackend, _ComputeNumpyBackend
from omnigibson.utils.backend_utils import _compute_backend as cb

from omnigibson.controllers import ControlType, ManipulationController
from omnigibson.controllers.joint_controller import JointController
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.utils.backend_utils import _ComputeBackend, _ComputeNumpyBackend, _ComputeTorchBackend
from omnigibson.utils.processing_utils import MovingAverageFilter
from omnigibson.utils.ui_utils import create_module_logger

@@ -220,7 +220,11 @@ def serialize(self, state):
return th.cat(
[
state_flat,
th.tensor([]) if self.control_filter is None else self.control_filter.serialize(state=state["control_filter"]),
(
th.tensor([])
if self.control_filter is None
else self.control_filter.serialize(state=state["control_filter"])
),
]
)

@@ -368,7 +372,10 @@ def command_dim(self):


import torch as th
cremebrule marked this conversation as resolved.
Show resolved Hide resolved

import omnigibson.utils.transform_utils as TT


@th.jit.script
def _compute_ik_qpos_torch(
q: th.Tensor,
@@ -400,7 +407,10 @@ def _compute_ik_qpos_torch(

import numpy as np
from numba import jit

import omnigibson.utils.transform_utils_np as NT


# Use numba since faster
@jit(nopython=True)
def _compute_ik_qpos_numpy(
@@ -432,4 +442,3 @@ def _compute_ik_qpos_numpy(
setattr(_ComputeBackend, "compute_ik_qpos", None)
setattr(_ComputeTorchBackend, "compute_ik_qpos", _compute_ik_qpos_torch)
setattr(_ComputeNumpyBackend, "compute_ik_qpos", _compute_ik_qpos_numpy)

12 changes: 7 additions & 5 deletions omnigibson/controllers/joint_controller.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
import math

from omnigibson.utils.backend_utils import _ComputeBackend, _ComputeTorchBackend, _ComputeNumpyBackend
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.controllers import (
ControlType,
GripperController,
@@ -10,6 +8,8 @@
ManipulationController,
)
from omnigibson.macros import create_module_macros
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.utils.backend_utils import _ComputeBackend, _ComputeNumpyBackend, _ComputeTorchBackend
from omnigibson.utils.python_utils import assert_valid_key
from omnigibson.utils.ui_utils import create_module_logger

@@ -286,6 +286,8 @@ def command_dim(self):


import torch as th


@th.compile
def _compute_joint_torques_torch(
u: th.Tensor,
@@ -296,9 +298,10 @@ def _compute_joint_torques_torch(
return mm[dof_idxs_mat] @ u



import numpy as np
from numba import jit


# Use numba since faster
@jit(nopython=True)
def numba_ix(arr, rows, cols):
@@ -316,7 +319,7 @@ def numba_ix(arr, rows, cols):
one_d_index = np.zeros(len(rows) * len(cols), dtype=np.int32)
for i, r in enumerate(rows):
start = i * len(cols)
one_d_index[start: start + len(cols)] = cols + arr.shape[1] * r
one_d_index[start : start + len(cols)] = cols + arr.shape[1] * r

arr_1d = arr.reshape((arr.shape[0] * arr.shape[1], 1))
slice_1d = np.take(arr_1d, one_d_index)
@@ -336,4 +339,3 @@ def _compute_joint_torques_numpy(
setattr(_ComputeBackend, "compute_joint_torques", None)
setattr(_ComputeTorchBackend, "compute_joint_torques", _compute_joint_torques_torch)
setattr(_ComputeNumpyBackend, "compute_joint_torques", _compute_joint_torques_numpy)

6 changes: 4 additions & 2 deletions omnigibson/controllers/multi_finger_gripper_controller.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.controllers import ControlType, GripperController, IsGraspingState
from omnigibson.macros import create_module_macros
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.utils.python_utils import assert_valid_key

VALID_MODES = {
@@ -254,7 +254,9 @@ def _update_grasping_state(self, control_dict):
dist_from_upper_limit = max_pos - finger_pos

# If either of the joint positions are not near the joint limits with some tolerance (m.POS_TOLERANCE)
valid_grasp_pos = dist_from_lower_limit.mean() > m.POS_TOLERANCE or dist_from_upper_limit.mean() > m.POS_TOLERANCE
valid_grasp_pos = (
dist_from_lower_limit.mean() > m.POS_TOLERANCE or dist_from_upper_limit.mean() > m.POS_TOLERANCE
)

# And the joint velocities are close to zero with some tolerance (m.VEL_TOLERANCE)
valid_grasp_vel = cb.all(cb.abs(finger_vel) < m.VEL_TOLERANCE)
2 changes: 1 addition & 1 deletion omnigibson/controllers/null_joint_controller.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.controllers import JointController
from omnigibson.utils.backend_utils import _compute_backend as cb


class NullJointController(JointController):
16 changes: 13 additions & 3 deletions omnigibson/controllers/osc_controller.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
import math

from omnigibson.utils.backend_utils import _ComputeBackend, _ComputeTorchBackend, _ComputeNumpyBackend
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.controllers import ControlType, ManipulationController
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.utils.backend_utils import _ComputeBackend, _ComputeNumpyBackend, _ComputeTorchBackend
from omnigibson.utils.processing_utils import MovingAverageFilter
from omnigibson.utils.python_utils import assert_valid_key
from omnigibson.utils.ui_utils import create_module_logger
@@ -412,7 +412,11 @@ def compute_control(self, goal_dict, control_dict):
ee_pos=cb.as_float32(ee_pos),
ee_mat=cb.as_float32(cb.T.quat2mat(ee_quat)),
ee_lin_vel=cb.as_float32(ee_vel[:3]),
ee_ang_vel_err=cb.as_float32(cb.T.quat2axisangle(cb.T.quat_multiply(cb.T.axisangle2quat(-ee_vel[3:]), cb.T.axisangle2quat(base_ang_vel)))),
ee_ang_vel_err=cb.as_float32(
cb.T.quat2axisangle(
cb.T.quat_multiply(cb.T.axisangle2quat(-ee_vel[3:]), cb.T.axisangle2quat(base_ang_vel))
)
),
cremebrule marked this conversation as resolved.
Show resolved Hide resolved
goal_pos=goal_dict["target_pos"],
goal_ori_mat=goal_dict["target_ori_mat"],
kp=kp,
@@ -486,7 +490,10 @@ def command_dim(self):


import torch as th

import omnigibson.utils.transform_utils as TT


@th.jit.script
def _compute_osc_torques_torch(
q: th.Tensor,
@@ -564,7 +571,10 @@ def _compute_osc_torques_torch(

import numpy as np
from numba import jit

import omnigibson.utils.transform_utils_np as NT


# Use numba since faster
@jit(nopython=True)
def _compute_osc_torques_numpy(
2 changes: 1 addition & 1 deletion omnigibson/objects/controllable_object.py
Original file line number Diff line number Diff line change
@@ -10,8 +10,8 @@
import omnigibson as og
from omnigibson.controllers import create_controller
from omnigibson.controllers.controller_base import ControlType
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.objects.object_base import BaseObject
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.utils.constants import JointType, PrimType
from omnigibson.utils.numpy_utils import NumpyTypes
from omnigibson.utils.python_utils import CachedFunctions, assert_valid_key, merge_nested_dicts
3 changes: 2 additions & 1 deletion omnigibson/robots/articulated_trunk_robot.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import torch as th
from functools import cached_property

import torch as th

from omnigibson.robots.manipulation_robot import ManipulationRobot
from omnigibson.utils.python_utils import classproperty
from omnigibson.utils.usd_utils import ControllableObjectViewAPI
2 changes: 1 addition & 1 deletion omnigibson/robots/behavior_robot.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import itertools
from functools import cached_property
import math
import os
from abc import ABC
from collections import OrderedDict
from functools import cached_property
from typing import Iterable, List, Literal, Tuple

import torch as th
10 changes: 7 additions & 3 deletions omnigibson/robots/holonomic_base_robot.py
Original file line number Diff line number Diff line change
@@ -6,10 +6,10 @@

import omnigibson as og
import omnigibson.lazy as lazy
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.macros import create_module_macros
from omnigibson.robots.locomotion_robot import LocomotionRobot
from omnigibson.robots.manipulation_robot import ManipulationRobot
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.utils.python_utils import classproperty
from omnigibson.utils.usd_utils import ControllableObjectViewAPI

@@ -254,7 +254,9 @@ def set_position_orientation(self, position=None, orientation=None, frame: Liter
joint_pos, joint_orn = cb.from_torch(joint_pos), cb.from_torch(joint_orn)
inv_joint_pos, inv_joint_orn = cb.T.mat2pose(cb.T.pose_inv(cb.T.pose2mat((joint_pos, joint_orn))))

relative_pos, relative_orn = cb.T.pose_transform(inv_joint_pos, inv_joint_orn, cb.from_torch(position), cb.from_torch(orientation))
relative_pos, relative_orn = cb.T.pose_transform(
inv_joint_pos, inv_joint_orn, cb.from_torch(position), cb.from_torch(orientation)
)
relative_rpy = cb.T.quat2euler(relative_orn)
self.joints["base_footprint_x_joint"].set_pos(relative_pos[0], drive=False)
self.joints["base_footprint_y_joint"].set_pos(relative_pos[1], drive=False)
@@ -321,7 +323,9 @@ def _postprocess_control(self, control, control_type):
global_pose = cur_pose @ local_pose
lin_vel_global, ang_vel_global = global_pose[0, :3, 3], global_pose[1, :3, 3]

u_vec[cb.from_torch(self.base_control_idx)] = cb.array([lin_vel_global[0], lin_vel_global[1], ang_vel_global[2]])
u_vec[cb.from_torch(self.base_control_idx)] = cb.array(
[lin_vel_global[0], lin_vel_global[1], ang_vel_global[2]]
)

return u_vec, u_type_vec

4 changes: 2 additions & 2 deletions omnigibson/robots/manipulation_robot.py
Original file line number Diff line number Diff line change
@@ -1,16 +1,15 @@
import math
import os
from abc import abstractmethod
from functools import cached_property
from collections import namedtuple
from functools import cached_property
from typing import Literal

import networkx as nx
import torch as th

import omnigibson as og
import omnigibson.lazy as lazy
from omnigibson.utils.backend_utils import _compute_backend as cb
import omnigibson.utils.transform_utils as T
from omnigibson.controllers import (
ControlType,
@@ -24,6 +23,7 @@
from omnigibson.macros import create_module_macros, gm
from omnigibson.object_states import ContactBodies
from omnigibson.robots.robot_base import BaseRobot
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.utils.constants import JointType, PrimType
from omnigibson.utils.geometry_utils import generate_points_in_volume_checker_function
from omnigibson.utils.python_utils import assert_valid_key, classproperty
2 changes: 1 addition & 1 deletion omnigibson/robots/tiago.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import math
import os
from typing import Literal
from functools import cached_property
from typing import Literal

import torch as th

5 changes: 2 additions & 3 deletions omnigibson/simulator.py
Original file line number Diff line number Diff line change
@@ -242,10 +242,9 @@ def _launch_app():

# Set compute backend
import omnigibson.utils.backend_utils as _backend_utils

_backend_utils._compute_backend.set_methods(
_backend_utils._ComputeNumpyBackend
if gm.USE_NUMPY_CONTROLLER_BACKEND
else _backend_utils._ComputeTorchBackend
_backend_utils._ComputeNumpyBackend if gm.USE_NUMPY_CONTROLLER_BACKEND else _backend_utils._ComputeTorchBackend
)

return app
3 changes: 2 additions & 1 deletion omnigibson/utils/backend_utils.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import torch as th
import numpy as np
import torch as th

import omnigibson.utils.transform_utils as TT
import omnigibson.utils.transform_utils_np as NT
from omnigibson.utils.python_utils import recursively_convert_from_torch
2 changes: 1 addition & 1 deletion omnigibson/utils/processing_utils.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import torch as th

from omnigibson.utils.python_utils import Serializable
from omnigibson.utils.backend_utils import _compute_backend as cb
from omnigibson.utils.python_utils import Serializable


class Filter(Serializable):
1 change: 1 addition & 0 deletions omnigibson/utils/python_utils.py
Original file line number Diff line number Diff line change
@@ -786,6 +786,7 @@ def recursively_convert_to_torch(state):
def recursively_convert_from_torch(state):
# For all the lists in state dict, convert from torch tensor -> numpy array
import numpy as np

for key, value in state.items():
if isinstance(value, dict):
state[key] = recursively_convert_from_torch(value)
1 change: 0 additions & 1 deletion omnigibson/utils/transform_utils.py
Original file line number Diff line number Diff line change
@@ -1402,4 +1402,3 @@ def orientation_error(desired, current):
error = 0.5 * (torch.linalg.cross(rc1, rd1) + torch.linalg.cross(rc2, rd2) + torch.linalg.cross(rc3, rd3))

return error.reshape(desired.shape[:-2] + (3,))

Loading