p04798526
/
OmniGibson
mirror of https://github.com/StanfordVL/OmniGibson.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #608 from StanfordVL/feat/controller-optimization 

Optimize Controller Logic
This commit is contained in:
Josiah Wong 2024-02-16 14:08:48 -08:00 committed by GitHub
parent f4a75fec7a 4388a7aa38
commit 191c821a1c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 63 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
omnigibson/objects/controllable_object.py
View File

@ -512,9 +512,9 @@ class ControllableObject(BaseObject):
fcns["joint_position"] = lambda: self.get_joint_positions(normalized=False)
fcns["joint_velocity"] = lambda: self.get_joint_velocities(normalized=False)
fcns["joint_effort"] = lambda: self.get_joint_efforts(normalized=False)
fcns["mass_matrix"] = self.get_mass_matrix
fcns["gravity_force"] = self.get_generalized_gravity_forces
fcns["cc_force"] = self.get_coriolis_and_centrifugal_forces
fcns["mass_matrix"] = lambda: self.get_mass_matrix(clone=False)
fcns["gravity_force"] = lambda: self.get_generalized_gravity_forces(clone=False)
fcns["cc_force"] = lambda: self.get_coriolis_and_centrifugal_forces(clone=False)
return fcns

6
omnigibson/objects/object_base.py
View File

@ -165,6 +165,12 @@ class BaseObject(EntityPrim, Registerable, metaclass=ABCMeta):
body1=f"{self._prim_path}/{self._root_link_name}",
)
# Delete n_fixed_joints cached property if it exists since the number of fixed joints has now changed
# See https://stackoverflow.com/questions/59899732/python-cached-property-how-to-delete and
# https://docs.python.org/3/library/functools.html#functools.cached_property
if "n_fixed_joints" in self.__dict__:
del self.n_fixed_joints
# Set visibility
if "visible" in self._load_config and self._load_config["visible"] is not None:
self.visible = self._load_config["visible"]

54
omnigibson/prims/entity_prim.py
View File

@ -1,5 +1,6 @@
import numpy as np
import networkx as nx
from functools import cached_property
import omnigibson as og
import omnigibson.lazy as lazy
@ -420,7 +421,7 @@ class EntityPrim(XFormPrim):
int: Number of joints owned by this articulation
"""
if self.initialized:
num = len(list(self._joints.keys()))
num = len(self._joints)
else:
# Manually iterate over all links and check for any joints that are not fixed joints!
num = 0
@ -433,7 +434,7 @@ class EntityPrim(XFormPrim):
num += 1
return num
@property
@cached_property
def n_fixed_joints(self):
"""
Returns:
@ -616,13 +617,12 @@ class EntityPrim(XFormPrim):
positions = self._denormalize_positions(positions=positions, indices=indices)
# Set the DOF states
if not drive:
if drive:
self._articulation_view.set_joint_position_targets(positions, joint_indices=indices)
else:
self._articulation_view.set_joint_positions(positions, joint_indices=indices)
BoundingBoxAPI.clear()
# Also set the target
self._articulation_view.set_joint_position_targets(positions, joint_indices=indices)
def set_joint_velocities(self, velocities, indices=None, normalized=False, drive=False):
"""
Set the joint velocities (both actual value and target values) in simulation. Note: only works if the simulator
@ -648,12 +648,11 @@ class EntityPrim(XFormPrim):
velocities = self._denormalize_velocities(velocities=velocities, indices=indices)
# Set the DOF states
if not drive:
if drive:
self._articulation_view.set_joint_velocity_targets(velocities, joint_indices=indices)
else:
self._articulation_view.set_joint_velocities(velocities, joint_indices=indices)
# Also set the target
self._articulation_view.set_joint_velocity_targets(velocities, joint_indices=indices)
def set_joint_efforts(self, efforts, indices=None, normalized=False):
"""
Set the joint efforts (both actual value and target values) in simulation. Note: only works if the simulator
@ -1247,49 +1246,64 @@ class EntityPrim(XFormPrim):
return aabb_lo, aabb_hi
def get_coriolis_and_centrifugal_forces(self):
def get_coriolis_and_centrifugal_forces(self, clone=True):
"""
Args:
clone (bool): Whether to clone the underlying tensor buffer or not
Returns:
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().reshape(self.n_dof)
return self._articulation_view.get_coriolis_and_centrifugal_forces(clone=clone).reshape(self.n_dof)
def get_generalized_gravity_forces(self):
def get_generalized_gravity_forces(self, clone=True):
"""
Args:
clone (bool): Whether to clone the underlying tensor buffer or not
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().reshape(self.n_dof)
return self._articulation_view.get_generalized_gravity_forces(clone=clone).reshape(self.n_dof)
def get_mass_matrix(self):
def get_mass_matrix(self, clone=True):
"""
Args:
clone (bool): Whether to clone the underlying tensor buffer or not
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().reshape(self.n_dof, self.n_dof)
return self._articulation_view.get_mass_matrices(clone=clone).reshape(self.n_dof, self.n_dof)
def get_jacobian(self):
def get_jacobian(self, clone=True):
"""
Args:
clone (bool): Whether to clone the underlying tensor buffer or not
Returns:
n-array: (N_links - 1 [+ 1], 6, N_dof [+ 6]) shaped per-link jacobian, if articulated. Note that the first
dimension is +1 and the final dimension is +6 if the entity does not have a fixed base
(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().squeeze(axis=0)
return self._articulation_view.get_jacobians(clone=clone).squeeze(axis=0)
def get_relative_jacobian(self):
def get_relative_jacobian(self, clone=True):
"""
Args:
clone (bool): Whether to clone the underlying tensor buffer or not
Returns:
n-array: (N_links - 1 [+ 1], 6, N_dof [+ 6]) shaped per-link relative jacobian, if articulated (expressed in
this entity's base frame). Note that the first dimension is +1 and the final dimension is +6 if the
entity does not have a fixed base (i.e.: there is an additional "floating" joint tying the robot to
the world frame)
"""
jac = self.get_jacobian()
jac = self.get_jacobian(clone=clone)
ori_t = T.quat2mat(self.get_orientation()).T.astype(np.float32)
tf = np.zeros((1, 6, 6), dtype=np.float32)
tf[:, :3, :3] = ori_t

2
omnigibson/robots/manipulation_robot.py
View File

@ -404,7 +404,7 @@ class ManipulationRobot(BaseRobot):
# not have a fixed base (i.e.: the 6DOF --> "floating" joint)
# see self.get_relative_jacobian() for more info
eef_link_idx = self._articulation_view.get_body_index(self.eef_links[arm].body_name)
fcns[f"eef_{arm}_jacobian_relative"] = lambda: self.get_relative_jacobian()[eef_link_idx, :, -self.n_joints:]
fcns[f"eef_{arm}_jacobian_relative"] = lambda: self.get_relative_jacobian(clone=False)[eef_link_idx, :, -self.n_joints:]
def _get_proprioception_dict(self):
dic = super()._get_proprioception_dict()

28
omnigibson/robots/robot_base.py
View File

@ -271,6 +271,22 @@ class BaseRobot(USDObject, ControllableObject, GymObservable):
"""
pass
def step(self):
# Skip this step if our articulation view is not valid
if self._articulation_view_direct is None or not self._articulation_view_direct.initialized:
return
# Before calling super, update the dummy robot's kinematic state based on this robot's kinematic state
# This is done prior to any state getter calls, since setting kinematic state results in physx backend
# having to re-fetch tensorized state.
# We do this so we have more optimal runtime performance
if self._use_dummy:
self._dummy.set_joint_positions(self.get_joint_positions())
self._dummy.set_joint_velocities(self.get_joint_velocities())
self._dummy.set_position_orientation(*self.get_position_orientation())
super().step()
def get_obs(self):
"""
Grabs all observations from the robot. This is keyword-mapped based on each observation modality
@ -471,16 +487,10 @@ class BaseRobot(USDObject, ControllableObject, GymObservable):
"""
return np.zeros(self.action_dim)
def get_generalized_gravity_forces(self):
def get_generalized_gravity_forces(self, clone=True):
# Override method based on whether we're using a dummy or not
if self._use_dummy:
# Update dummy pose and calculate values
self._dummy.set_joint_positions(self.get_joint_positions())
self._dummy.set_joint_velocities(self.get_joint_velocities())
self._dummy.set_position_orientation(*self.get_position_orientation())
return self._dummy.get_generalized_gravity_forces()
else:
return super().get_generalized_gravity_forces()
return self._dummy.get_generalized_gravity_forces(clone=clone) \
if self._use_dummy else super().get_generalized_gravity_forces(clone=clone)
@property
def sensors(self):