diff --git a/omnigibson/__init__.py b/omnigibson/__init__.py index b25689a46..9a5557ae1 100644 --- a/omnigibson/__init__.py +++ b/omnigibson/__init__.py @@ -80,6 +80,11 @@ def create_app(): config_kwargs["active_gpu"] = gpu_id config_kwargs["physics_gpu"] = gpu_id app = SimulationApp(config_kwargs) + + # Omni overrides the global logger to be DEBUG, which is very annoying, so we re-override it to the default WARN + # TODO: Remove this once omniverse fixes it + logging.getLogger().setLevel(logging.WARNING) + import omni # Enable additional extensions we need diff --git a/omnigibson/controllers/__init__.py b/omnigibson/controllers/__init__.py index be53954fb..7cdcb65b0 100644 --- a/omnigibson/controllers/__init__.py +++ b/omnigibson/controllers/__init__.py @@ -9,6 +9,7 @@ GripperController, ) from omnigibson.controllers.dd_controller import DifferentialDriveController +from omnigibson.controllers.osc_controller import OperationalSpaceController from omnigibson.controllers.ik_controller import InverseKinematicsController from omnigibson.controllers.joint_controller import JointController from omnigibson.controllers.multi_finger_gripper_controller import MultiFingerGripperController diff --git a/omnigibson/controllers/osc_controller.py b/omnigibson/controllers/osc_controller.py new file mode 100644 index 000000000..450f11373 --- /dev/null +++ b/omnigibson/controllers/osc_controller.py @@ -0,0 +1,532 @@ +import numpy as np +from numba import jit + +import omnigibson.utils.transform_utils as T +from omnigibson.controllers import ControlType, ManipulationController +from omnigibson.utils.control_utils import orientation_error +from omnigibson.utils.processing_utils import MovingAverageFilter +from omnigibson.utils.python_utils import nums2array, assert_valid_key +from omnigibson.utils.ui_utils import create_module_logger + +# Create module logger +log = create_module_logger(module_name=__name__) + +# Different modes +OSC_MODE_COMMAND_DIMS = { + "absolute_pose": 6, # 6DOF (x,y,z,ax,ay,az) control of pose, whether both position and orientation is given in absolute coordinates + "pose_absolute_ori": 6, # 6DOF (dx,dy,dz,ax,ay,az) control over pose, where the orientation is given in absolute axis-angle coordinates + "pose_delta_ori": 6, # 6DOF (dx,dy,dz,dax,day,daz) control over pose + "position_fixed_ori": 3, # 3DOF (dx,dy,dz) control over position, with orientation commands being kept as fixed initial absolute orientation + "position_compliant_ori": 3, # 3DOF (dx,dy,dz) control over position, with orientation commands automatically being sent as 0s (so can drift over time) +} +OSC_MODES = set(OSC_MODE_COMMAND_DIMS.keys()) + + +class OperationalSpaceController(ManipulationController): + """ + Controller class to convert (delta or absolute) EEF commands into joint efforts using Operational Space Control + + This controller expects 6DOF delta commands (dx, dy, dz, dax, day, daz), where the delta orientation + commands are in axis-angle form, and outputs low-level torque commands. + + Gains may also be considered part of the action space as well. In this case, the action space would be: + ( + dx, dy, dz, dax, day, daz <-- 6DOF delta eef commands + [, kpx, kpy, kpz, kpax, kpay, kpaz] <-- kp gains + [, drx dry, drz, drax, dray, draz] <-- damping ratio gains + [, kpnx, kpny, kpnz, kpnax, kpnay, kpnaz] <-- kp null gains + ) + + Note that in this case, we ASSUME that the inputted gains are normalized to be in the range [-1, 1], and will + be mapped appropriately to their respective ranges, as defined by XX_limits + + Alternatively, parameters (in this case, kp or damping_ratio) can either be set during initialization or provided + from an external source; if the latter, the control_dict should include the respective parameter(s) as + a part of its keys + + Each controller step consists of the following: + 1. Clip + Scale inputted command according to @command_input_limits and @command_output_limits + 2. Run OSC to back out joint efforts for a desired task frame command + 3. Clips the resulting command by the motor (effort) limits + """ + + def __init__( + self, + task_name, + control_freq, + default_joint_pos, + control_limits, + dof_idx, + command_input_limits="default", + command_output_limits=((-0.2, -0.2, -0.2, -0.5, -0.5, -0.5), (0.2, 0.2, 0.2, 0.5, 0.5, 0.5)), + kp=150.0, + kp_limits=(10.0, 300.), + damping_ratio=1.0, + damping_ratio_limits=(0.0, 2.0), + kp_null=10.0, + kp_null_limits=(0.0, 50.0), + mode="pose_delta_ori", + decouple_pos_ori=False, + workspace_pose_limiter=None, + ): + """ + Args: + task_name (str): name assigned to this task frame for computing OSC control. During control calculations, + the inputted control_dict should include entries named <@task_name>_pos_relative and + <@task_name>_quat_relative. See self._command_to_control() for what these values should entail. + control_freq (int): controller loop frequency + default_joint_pos (Array[float]): default joint positions, used as part of nullspace controller in IK. + Note that this should correspond to ALL the joints; the exact indices will be extracted via @dof_idx + control_limits (Dict[str, Tuple[Array[float], Array[float]]]): The min/max limits to the outputted + control signal. Should specify per-dof type limits, i.e.: + + "position": [[min], [max]] + "velocity": [[min], [max]] + "effort": [[min], [max]] + "has_limit": [...bool...] + + Values outside of this range will be clipped, if the corresponding joint index in has_limit is True. + dof_idx (Array[int]): specific dof indices controlled by this robot. Used for inferring + controller-relevant values during control computations + command_input_limits (None or "default" or Tuple[float, float] or Tuple[Array[float], Array[float]]): + if set, is the min/max acceptable inputted command. Values outside this range will be clipped. + If None, no clipping will be used. If "default", range will be set to (-1, 1) + command_output_limits (None or "default" or Tuple[float, float] or Tuple[Array[float], Array[float]]): + if set, is the min/max scaled command. If both this value and @command_input_limits is not None, + then all inputted command values will be scaled from the input range to the output range. + If either is None, no scaling will be used. If "default", then this range will automatically be set + to the @control_limits entry corresponding to self.control_type + kp (None, int, float, or array): Gain values to apply to 6DOF error. + If None, will be variable (part of action space) + kp_limits (2-array): (min, max) values of kp + damping_ratio (None, int, float, or array): Damping ratio to apply to 6DOF error controller gain + If None, will be variable (part of action space) + damping_ratio_limits (2-array): (min, max) values of damping ratio + kp_null (None, int, float, or array): Gain applied when calculating null torques + If None, will be variable (part of action space) + kp_null_limits (2-array): (min, max) values of kp_null + mode (str): mode to use when computing IK. In all cases, position commands are 3DOF delta (dx,dy,dz) + cartesian values, relative to the robot base frame. Valid options are: + - "pose_absolute_ori": 6DOF (dx,dy,dz,ax,ay,az) control over pose, + where the orientation is given in absolute axis-angle coordinates + - "pose_delta_ori": 6DOF (dx,dy,dz,dax,day,daz) control over pose + - "position_fixed_ori": 3DOF (dx,dy,dz) control over position, + with orientation commands being kept as fixed initial absolute orientation + - "position_compliant_ori": 3DOF (dx,dy,dz) control over position, + with orientation commands automatically being sent as 0s (so can drift over time) + decouple_pos_ori (bool): Whether to decouple position and orientation control or not + workspace_pose_limiter (None or function): if specified, callback method that should clip absolute + target (x,y,z) cartesian position and absolute quaternion orientation (x,y,z,w) to a specific workspace + range (i.e.: this can be unique to each robot, and implemented by each embodiment). + Function signature should be: + + def limiter(command_pos: Array[float], command_quat: Array[float], control_dict: Dict[str, Any]) --> Tuple[Array[float], Array[float]] + + where pos_command is (x,y,z) cartesian position values, command_quat is (x,y,z,w) quarternion orientation + values, and the returned tuple is the processed (pos, quat) command. + """ + # Store arguments + control_dim = len(dof_idx) + + # Store gains + self.kp = nums2array(nums=kp, dim=6, dtype=np.float32) if kp is not None else None + self.damping_ratio = damping_ratio + self.kp_null = nums2array(nums=kp_null, dim=control_dim, dtype=np.float32) if kp_null is not None else None + self.kd_null = 2 * np.sqrt(self.kp_null) if kp_null is not None else None # critically damped + self.kp_limits = np.array(kp_limits, dtype=np.float32) + self.damping_ratio_limits = np.array(damping_ratio_limits, dtype=np.float32) + self.kp_null_limits = np.array(kp_null_limits, dtype=np.float32) + + # Store settings for whether we're learning gains or not + self.variable_kp = self.kp is None + self.variable_damping_ratio = self.damping_ratio is None + self.variable_kp_null = self.kp_null is None + + # If the mode is set as absolute orientation and using default config, + # change input and output limits accordingly. + # By default, the input limits are set as 1, so we modify this to have a correct range. + # The output orientation limits are also set to be values assuming delta commands, so those are updated too + assert_valid_key(key=mode, valid_keys=OSC_MODES, name="OSC mode") + self.mode = mode + if self.mode == "pose_absolute_ori": + if command_input_limits is not None: + if command_input_limits == "default": + command_input_limits = [ + [-1.0, -1.0, -1.0, -np.pi, -np.pi, -np.pi], + [1.0, 1.0, 1.0, np.pi, np.pi, np.pi], + ] + else: + command_input_limits[0][3:] = -np.pi + command_input_limits[1][3:] = np.pi + if command_output_limits is not None: + if command_output_limits == "default": + command_output_limits = [ + [-1.0, -1.0, -1.0, -np.pi, -np.pi, -np.pi], + [1.0, 1.0, 1.0, np.pi, np.pi, np.pi], + ] + else: + command_output_limits[0][3:] = -np.pi + command_output_limits[1][3:] = np.pi + + is_input_limits_numeric = not (command_input_limits is None or isinstance(command_input_limits, str)) + is_output_limits_numeric = not (command_output_limits is None or isinstance(command_output_limits, str)) + command_input_limits = [nums2array(lim, dim=6, dtype=np.float32) for lim in command_input_limits] if is_input_limits_numeric else command_input_limits + command_output_limits = [nums2array(lim, dim=6, dtype=np.float32) for lim in command_output_limits] if is_output_limits_numeric else command_output_limits + + # Modify input / output scaling based on whether we expect gains to be part of the action space + self._command_dim = OSC_MODE_COMMAND_DIMS[self.mode] + for variable_gain, gain_limits, dim in zip( + (self.variable_kp, self.variable_damping_ratio, self.variable_kp_null), + (self.kp_limits, self.damping_ratio_limits, self.kp_null_limits), + (6, 6, control_dim), + ): + if variable_gain: + # Add this to input / output limits + if is_input_limits_numeric: + command_input_limits = [np.concatenate([lim, nums2array(nums=val, dim=dim, dtype=np.float32)]) for lim, val in zip(command_input_limits, (-1, 1))] + if is_output_limits_numeric: + command_output_limits = [np.concatenate([lim, nums2array(nums=val, dim=dim, dtype=np.float32)]) for lim, val in zip(command_output_limits, gain_limits)] + # Update command dim + self._command_dim += dim + + # Other values + self.decouple_pos_ori = decouple_pos_ori + self.workspace_pose_limiter = workspace_pose_limiter + self.task_name = task_name + self.default_joint_pos = default_joint_pos[dof_idx].astype(np.float32) + + # Other variables that will be filled in at runtime + self.goal_pos = None + self.goal_ori_mat = None + self._quat_target = None + + # Run super init + super().__init__( + control_freq=control_freq, + control_limits=control_limits, + dof_idx=dof_idx, + command_input_limits=command_input_limits, + command_output_limits=command_output_limits, + ) + + def reset(self): + # Clear internal variables + self.goal_pos = None + self.goal_ori_mat = None + self._quat_target = None + self._clear_variable_gains() + + @property + def state_size(self): + # Add 4 for internal quat target + # TODO: Store goal pose, goal ori mat, gains, etc. + return super().state_size + 4 + + def _clear_variable_gains(self): + """ + Helper function to clear any gains that are variable and considered part of actions + """ + if self.variable_kp: + self.kp = None + if self.variable_damping_ratio: + self.damping_ratio = None + if self.variable_kp_null: + self.kp_null = None + self.kd_null = None + + def _update_variable_gains(self, gains): + """ + Helper function to update any gains that are variable and considered part of actions + + Args: + gains (n-array): array where n dim is parsed based on which gains are being learned + """ + idx = 0 + if self.variable_kp: + self.kp = gains[:, idx:idx + 6].astype(np.float32) + idx += 6 + if self.variable_damping_ratio: + self.damping_ratio = gains[:, idx:idx + 6].astype(np.float32) + idx += 6 + if self.variable_kp_null: + self.kp_null = gains[:, idx:idx + self.control_dim].astype(np.float32) + self.kd_null = 2 * np.sqrt(self.kp_null) # critically damped + idx += self.control_dim + + def _dump_state(self): + # Run super first + state = super()._dump_state() + + # Add internal quaternion target and filter state + state["quat_target"] = self._quat_target + + return state + + def _load_state(self, state): + # Run super first + super()._load_state(state=state) + + # Load relevant info for this controller + self._quat_target = state["quat_target"] + + def _serialize(self, state): + # Run super first + state_flat = super()._serialize(state=state) + + # Serialize state for this controller + return np.concatenate([ + state_flat, + np.zeros(4) if state["quat_target"] is None else state["quat_target"], # Encode None as zeros for consistent serialization size + ]).astype(float) + + def _deserialize(self, state): + # Run super first + state_dict, idx = super()._deserialize(state=state) + + # Deserialize state for this controller + state_dict["quat_target"] = None if np.all(state[idx: idx + 4] == 0.0) else state[idx: idx + 4] + + return state_dict, idx + 4 + + def update_goal(self, control_dict, target_pos, target_quat, gains=None): + """ + Updates the internal goal (ee pos and ee ori mat) based on the inputted delta command + + Args: + control_dict (Dict[str, Any]): dictionary that should include any relevant keyword-mapped + states necessary for controller computation. Must include the following keys: + joint_position: Array of current joint positions + <@self.task_name>_pos_relative: (x,y,z) relative cartesian position of the desired task frame to + control, computed in its local frame (e.g.: robot base frame) + <@self.task_name>_quat_relative: (x,y,z,w) relative quaternion orientation of the desired task + frame to control, computed in its local frame (e.g.: robot base frame) + <@self.task_name>_lin_vel_relative: (x,y,z) relative linear velocity of the desired task frame to + control, computed in its local frame (e.g.: robot base frame) + <@self.task_name>_ang_vel_relative: (ax, ay, az) relative angular velocity of the desired task + frame to control, computed in its local frame (e.g.: robot base frame) + + target_pos (3-array): (x,y,z) desired position of the target frame with respect to the robot frame. + These should be the DE-NORMALIZED commands (i.e.: already processed wrt to the input / output command + limits) + target_quat (4-array): (x,y,z,w) desired quaternion orientation of the target frame with respect to the + robot frame. These should be the DE-NORMALIZED commands (i.e.: already processed wrt to the input + / output command limits) + gains (n-array): If specified, gains to update internally (will only be used if variable_kp / etc. is used) + """ + # Set goals + self.goal_pos = target_pos.astype(np.float32) + self.goal_ori_mat = T.quat2mat(target_quat).astype(np.float32) + if gains is not None: + self._update_variable_gains(gains=gains) + + def compute_control(self, control_dict): + """ + Computes low-level torque controls using internal eef goal pos / ori. + + Args: + control_dict (Dict[str, Any]): dictionary that should include any relevant keyword-mapped + states necessary for controller computation. Must include the following keys: + joint_position: Array of current joint positions + joint_velocity: Array of current joint velocities + mass_matrix: (N_dof, N_dof) Current mass matrix + <@self.task_name>_jacobian_relative: (6, N_dof) Current jacobian matrix for desired task frame + <@self.task_name>_pos_relative: (x,y,z) relative cartesian position of the desired task frame to + control, computed in its local frame (e.g.: robot base frame) + <@self.task_name>_quat_relative: (x,y,z,w) relative quaternion orientation of the desired task + frame to control, computed in its local frame (e.g.: robot base frame) + <@self.task_name>_lin_vel_relative: (x,y,z) relative linear velocity of the desired task frame to + control, computed in its local frame (e.g.: robot base frame) + <@self.task_name>_ang_vel_relative: (ax, ay, az) relative angular velocity of the desired task + frame to control, computed in its local frame (e.g.: robot base frame) + + control_dict (dict): Dictionary of state tensors including relevant info for controller computation + + Returns: + n-array: low-level effort control actions, NOT post-processed + """ + # TODO: Update to possibly grab parameters from dict + # For now, always use internal values + kp = self.kp + damping_ratio = self.damping_ratio + kd = 2 * np.sqrt(kp) * damping_ratio + + # Extract relevant values from the control dict + dof_idxs_mat = tuple(np.meshgrid(self.dof_idx, self.dof_idx)) + q = control_dict["joint_position"][self.dof_idx] + qd = control_dict["joint_velocity"][self.dof_idx] + mm = control_dict["mass_matrix"][dof_idxs_mat] + j_eef = control_dict[f"{self.task_name}_jacobian_relative"][:, self.dof_idx] + ee_pos = control_dict[f"{self.task_name}_pos_relative"] + ee_quat = control_dict[f"{self.task_name}_quat_relative"] + ee_vel = np.concatenate([control_dict[f"{self.task_name}_lin_vel_relative"], control_dict[f"{self.task_name}_ang_vel_relative"]]) + + # Calculate torques + u = _compute_osc_torques( + q=q, + qd=qd, + mm=mm, + j_eef=j_eef, + ee_pos=ee_pos.astype(np.float32), + ee_mat=T.quat2mat(ee_quat).astype(np.float32), + ee_vel=ee_vel.astype(np.float32), + goal_pos=self.goal_pos, + goal_ori_mat=self.goal_ori_mat, + kp=kp, + kd=kd, + kp_null=self.kp_null, + kd_null=self.kd_null, + rest_qpos=self.default_joint_pos, + control_dim=self.control_dim, + decouple_pos_ori=self.decouple_pos_ori, + ) + + # Return the control torques + return u + + def _command_to_control(self, command, control_dict): + """ + Converts the (already preprocessed) inputted @command into deployable (non-clipped!) joint control signal. + This processes the command based on self.mode, possibly clips the command based on self.workspace_pose_limiter, + + Args: + command (Array[float]): desired (already preprocessed) command to convert into control signals + Is one of: + (dx,dy,dz,...) - desired delta cartesian position + (dx,dy,dz,dax,day,daz,...) - desired delta cartesian position and delta axis-angle orientation + (dx,dy,dz,ax,ay,az,...) - desired delta cartesian position and global axis-angle orientation + where ... represents potential additional gain values specified + control_dict (Dict[str, Any]): dictionary that should include any relevant keyword-mapped + states necessary for controller computation. Must include the following keys: + joint_position: Array of current joint positions + joint_velocity: Array of current joint velocities + mass_matrix: (N_dof, N_dof) Current mass matrix + gravity_force: (N_dof,) Gravity compensation efforts to offset for + cc_force: (N_dof,) Coriolis and centrifugal efforts to offset for + <@self.task_name>_pos_relative: (x,y,z) relative cartesian position of the desired task frame to + control, computed in its local frame (e.g.: robot base frame) + <@self.task_name>_quat_relative: (x,y,z,w) relative quaternion orientation of the desired task + frame to control, computed in its local frame (e.g.: robot base frame) + <@self.task_name>_lin_vel_relative: (x,y,z) relative linear velocity of the desired task frame to + control, computed in its local frame (e.g.: robot base frame) + <@self.task_name>_ang_vel_relative: (ax, ay, az) relative angular velocity of the desired task + frame to control, computed in its local frame (e.g.: robot base frame) + + Returns: + Array[float]: outputted (non-clipped!) velocity control signal to deploy + """ + # Grab important info from control dict + pos_relative = np.array(control_dict[f"{self.task_name}_pos_relative"]) + quat_relative = np.array(control_dict[f"{self.task_name}_quat_relative"]) + + # The first three values of the command are always the (delta) position, convert to absolute values if needed + if self.mode == "absolute_pose": + target_pos = command[:3] + else: + dpos = command[:3] + target_pos = pos_relative + dpos + + # Compute orientation + if self.mode == "position_fixed_ori": + # We need to grab the current robot orientation as the commanded orientation if there is none saved + if self._quat_target is None: + self._quat_target = quat_relative.astype(np.float32) + target_quat = self._quat_target + elif self.mode == "position_compliant_ori": + # Target quat is simply the current robot orientation + target_quat = quat_relative + elif self.mode == "pose_absolute_ori" or self.mode == "absolute_pose": + # Received "delta" ori is in fact the desired absolute orientation + target_quat = T.axisangle2quat(command[3:6]) + else: # pose_delta_ori control + # Grab dori and compute target ori + dori = T.quat2mat(T.axisangle2quat(command[3:6])) + target_quat = T.mat2quat(dori @ T.quat2mat(quat_relative)) + + # Possibly limit to workspace if specified + if self.workspace_pose_limiter is not None: + target_pos, target_quat = self.workspace_pose_limiter(target_pos, target_quat, control_dict) + + # Update goal + gains = command[OSC_MODE_COMMAND_DIMS[self.mode]:] + self.update_goal(control_dict=control_dict, target_pos=target_pos, target_quat=target_quat, gains=gains) + + # Calculate and return OSC-backed out joint efforts + u = self.compute_control(control_dict=control_dict).flatten() + + # Apply gravity compensation from the control dict + u += control_dict["gravity_force"][self.dof_idx] + control_dict["cc_force"][self.dof_idx] + + # Return these commanded velocities (this only includes the relevant dof idx) + return u + + @property + def control_type(self): + return ControlType.EFFORT + + @property + def command_dim(self): + return self._command_dim + + +# Use numba since faster +@jit(nopython=True) +def _compute_osc_torques( + q, + qd, + mm, + j_eef, + ee_pos, + ee_mat, + ee_vel, + goal_pos, + goal_ori_mat, + kp, + kd, + kp_null, + kd_null, + rest_qpos, + control_dim, + decouple_pos_ori, +): + # Compute the inverse + mm_inv = np.linalg.inv(mm) + + # Calculate error + pos_err = goal_pos - ee_pos + ori_err = orientation_error(goal_ori_mat, ee_mat).astype(np.float32) + err = np.concatenate((pos_err, ori_err)) + + # Determine desired wrench + err = np.expand_dims(kp * err - kd * ee_vel, axis=-1) + m_eef_inv = j_eef @ mm_inv @ j_eef.T + m_eef = np.linalg.inv(m_eef_inv) + + if decouple_pos_ori: + # # More efficient, but numba doesn't support 3D tensor operations yet + # j_eef_batch = j_eef.reshape(2, 3, -1) + # m_eef_pose_inv = np.matmul(np.matmul(j_eef_batch, np.expand_dims(mm_inv, axis=0)), np.transpose(j_eef_batch, (0, 2, 1))) + # m_eef_pose = np.linalg.inv(m_eef_pose_inv) # Shape (2, 3, 3) + # wrench = np.matmul(m_eef_pose, err.reshape(2, 3, 1)).flatten() + m_eef_pos_inv = j_eef[:3, :] @ mm_inv @ j_eef[:3, :].T + m_eef_ori_inv = j_eef[3:, :] @ mm_inv @ j_eef[3:, :].T + m_eef_pos = np.linalg.inv(m_eef_pos_inv) + m_eef_ori = np.linalg.inv(m_eef_ori_inv) + wrench_pos = m_eef_pos @ err[:3, :] + wrench_ori = m_eef_ori @ err[3:, :] + wrench = np.concatenate((wrench_pos, wrench_ori)) + else: + wrench = m_eef @ err + + # Compute OSC torques + u = j_eef.T @ wrench + + # Nullspace control torques `u_null` prevents large changes in joint configuration + # They are added into the nullspace of OSC so that the end effector orientation remains constant + # roboticsproceedings.org/rss07/p31.pdf + if rest_qpos is not None: + j_eef_inv = m_eef @ j_eef @ mm_inv + u_null = kd_null * -qd + kp_null * ((rest_qpos - q + np.pi) % (2 * np.pi) - np.pi) + u_null = mm @ np.expand_dims(u_null, axis=-1).astype(np.float32) + u += (np.eye(control_dim, dtype=np.float32) - j_eef.T @ j_eef_inv) @ u_null + + return u diff --git a/omnigibson/envs/__init__.py b/omnigibson/envs/__init__.py index 8cc3c969a..a93206368 100644 --- a/omnigibson/envs/__init__.py +++ b/omnigibson/envs/__init__.py @@ -1 +1,2 @@ from omnigibson.envs.env_base import Environment +from omnigibson.envs.env_wrapper import EnvironmentWrapper, create_wrapper, REGISTERED_ENV_WRAPPERS diff --git a/omnigibson/envs/env_base.py b/omnigibson/envs/env_base.py index 8a68e965d..9b5034e2f 100644 --- a/omnigibson/envs/env_base.py +++ b/omnigibson/envs/env_base.py @@ -1,5 +1,6 @@ import gym import numpy as np +from copy import deepcopy import omnigibson as og from omnigibson.objects import REGISTERED_OBJECTS @@ -51,7 +52,9 @@ def __init__( self._automatic_reset = automatic_reset self._flatten_action_space = flatten_action_space self._flatten_obs_space = flatten_obs_space + self.physics_timestep = physics_timestep self.action_timestep = action_timestep + self.device = device # Initialize other placeholders that will be filled in later self._initial_pos_z_offset = None # how high to offset object placement to account for one action step of dropping @@ -76,13 +79,7 @@ def __init__( self._scene_graph_builder = None if "scene_graph" in self.config and self.config["scene_graph"] is not None: self._scene_graph_builder = SceneGraphBuilder(**self.config["scene_graph"]) - - # Set the simulator settings - og.sim.set_simulation_dt(physics_dt=physics_timestep, rendering_dt=action_timestep) - og.sim.viewer_width = self.render_config["viewer_width"] - og.sim.viewer_height = self.render_config["viewer_height"] - og.sim.device = device - + # Load this environment self.load() @@ -159,10 +156,22 @@ def _load_variables(self): drop_distance = 0.5 * 9.8 * (self.action_timestep ** 2) assert drop_distance < self._initial_pos_z_offset, "initial_pos_z_offset is too small for collision checking" - def _load_task(self): + def _load_task(self, task_config=None): """ Load task + + Args: + task_confg (None or dict): If specified, custom task configuration to use. Otherwise, will use + self.task_config. Note that if a custom task configuration is specified, the internal task config + will be updated as well """ + # Update internal config if specified + if task_config is not None: + # Copy task config, in case self.task_config and task_config are the same! + task_config = deepcopy(task_config) + self.task_config.clear() + self.task_config.update(task_config) + # Sanity check task to make sure it's valid task_type = self.task_config["type"] assert_valid_key(key=task_type, valid_keys=REGISTERED_TASKS, name="task type") @@ -194,6 +203,13 @@ def _load_scene(self): cls_type_descriptor="scene", ) og.sim.import_scene(scene) + + # Set the simulator settings + og.sim.set_simulation_dt(physics_dt=self.physics_timestep, rendering_dt=self.action_timestep) + og.sim.viewer_width = self.render_config["viewer_width"] + og.sim.viewer_height = self.render_config["viewer_height"] + og.sim.device = self.device + assert og.sim.is_stopped(), "Simulator must be stopped after loading scene!" def _load_robots(self): @@ -334,6 +350,31 @@ def load(self): # Denote that the scene is loaded self._loaded = True + def update_task(self, task_config): + """ + Updates the internal task using @task_config. NOTE: This will internally reset the environment as well! + + Args: + task_config (dict): Task configuration for updating the new task + """ + # Make sure sim is playing + assert og.sim.is_playing(), "Update task should occur while sim is playing!" + + # Denote scene as not loaded yet + self._loaded = False + og.sim.stop() + self._load_task(task_config=task_config) + og.sim.play() + self.reset() + + # Load obs / action spaces + self.load_observation_space() + self._load_action_space() + + # Scene is now loaded again + self._loaded = True + + def close(self): """ Clean up the environment and shut down the simulation. @@ -471,7 +512,7 @@ def reset(self): # Grab and return observations obs = self.get_obs() - if self.observation_space is not None and not self.observation_space.contains(obs): + if self._loaded and not self.observation_space.contains(obs): # Flatten obs, and print out all keys and values log.error("OBSERVATION SPACE:") for key, value in recursively_generate_flat_dict(dic=self.observation_space).items(): @@ -571,6 +612,14 @@ def task_config(self): """ return self.config["task"] + @property + def wrapper_config(self): + """ + Returns: + dict: Wrapper-specific configuration kwargs + """ + return self.config["wrapper"] + @property def default_config(self): """ @@ -612,5 +661,10 @@ def default_config(self): # Task kwargs "task": { "type": "DummyTask", - } + }, + + # Wrapper kwargs + "wrapper": { + "type": None, + }, } diff --git a/omnigibson/envs/env_wrapper.py b/omnigibson/envs/env_wrapper.py index e49a96fa7..0561d6da8 100644 --- a/omnigibson/envs/env_wrapper.py +++ b/omnigibson/envs/env_wrapper.py @@ -1,7 +1,32 @@ from omnigibson.utils.python_utils import Wrapper +from omnigibson.utils.python_utils import Registerable, classproperty, create_class_from_registry_and_config +from omnigibson.utils.ui_utils import create_module_logger +from copy import deepcopy +# Global dicts that will contain mappings +REGISTERED_ENV_WRAPPERS = dict() -class EnvironmentWrapper(Wrapper): +# Create module logger +log = create_module_logger(module_name=__name__) + + +def create_wrapper(env): + """ + Wraps environment @env with wrapper defined by env.wrapper_config + """ + wrapper_cfg = deepcopy(env.wrapper_config) + wrapper_type = wrapper_cfg.pop("type") + wrapper_cfg["env"] = env + + return create_class_from_registry_and_config( + cls_name=wrapper_type, + cls_registry=REGISTERED_ENV_WRAPPERS, + cfg=wrapper_cfg, + cls_type_descriptor="wrapper", + ) + + +class EnvironmentWrapper(Wrapper, Registerable): """ Base class for all environment wrappers in OmniGibson. In general, reset(), step(), and observation_spec() should be overwritten @@ -50,3 +75,15 @@ def observation_spec(self): """ return self.env.observation_spec() + @classproperty + def _do_not_register_classes(cls): + # Don't register this class since it's an abstract template + classes = super()._do_not_register_classes + classes.add("EnvironmentWrapper") + return classes + + @classproperty + def _cls_registry(cls): + # Global robot registry + global REGISTERED_ENV_WRAPPERS + return REGISTERED_ENV_WRAPPERS diff --git a/omnigibson/objects/controllable_object.py b/omnigibson/objects/controllable_object.py index 5b51281ce..422fa0348 100644 --- a/omnigibson/objects/controllable_object.py +++ b/omnigibson/objects/controllable_object.py @@ -436,8 +436,8 @@ def deploy_control(self, control, control_type, indices=None, normalized=False): elif ctrl_type == ControlType.POSITION: joint.set_pos(ctrl, normalized=norm, drive=True) elif ctrl_type == ControlType.NONE: - # Do nothing - pass + # Set zero efforts + joint.set_effort(0, normalized=False) else: raise ValueError("Invalid control type specified: {}".format(ctrl_type)) @@ -457,6 +457,9 @@ def get_control_dict(self): - joint_effort: (n_dof,) joint efforts - root_pos: (3,) (x,y,z) global cartesian position of the object's root link - root_quat: (4,) (x,y,z,w) global cartesian orientation of ths object's root link + - mass_matrix: (n_dof, n_dof) mass matrix + - gravity_force: (n_dof,) per-joint generalized gravity forces + - cc_force: (n_dof,) per-joint centripetal and centrifugal forces """ pos, ori = self.get_position_orientation() return dict( @@ -465,6 +468,9 @@ def get_control_dict(self): joint_effort=self.get_joint_efforts(normalized=False), root_pos=pos, root_quat=ori, + mass_matrix=self.get_mass_matrix(), + gravity_force=self.get_generalized_gravity_forces(), + cc_force=self.get_coriolis_and_centrifugal_forces(), ) def dump_action(self): diff --git a/omnigibson/objects/light_object.py b/omnigibson/objects/light_object.py index 1805c483e..1835d1aa6 100644 --- a/omnigibson/objects/light_object.py +++ b/omnigibson/objects/light_object.py @@ -161,8 +161,7 @@ def radius(self): Returns: float: radius for this light """ - return self._light_link.get_attribute( - "inputs:radius" if meets_minimum_isaac_version("2023.0.0") else "radius") + return self._light_link.get_attribute("inputs:radius" if meets_minimum_isaac_version("2023.0.0") else "radius") @radius.setter def radius(self, radius): @@ -172,9 +171,7 @@ def radius(self, radius): Args: radius (float): radius to set """ - self._light_link.set_attribute( - "inputs:radius" if meets_minimum_isaac_version("2023.0.0") else "radius", - radius) + self._light_link.set_attribute("inputs:radius" if meets_minimum_isaac_version("2023.0.0") else "radius", radius) @property def intensity(self): diff --git a/omnigibson/prims/entity_prim.py b/omnigibson/prims/entity_prim.py index 14793c37d..9a959ee0c 100644 --- a/omnigibson/prims/entity_prim.py +++ b/omnigibson/prims/entity_prim.py @@ -59,6 +59,7 @@ def __init__( self._joints = None self._materials = None self._visual_only = None + self._physics_view = None # This needs to be initialized to be used for _load() of PrimitiveObject self._prim_type = load_config["prim_type"] if load_config is not None and "prim_type" in load_config else PrimType.RIGID @@ -94,6 +95,10 @@ def _initialize(self): # Update joint information self.update_joints() + # Construct physics view + if self.articulated: + self._physics_view = og.sim.physics_sim_view.create_articulation_view(self.articulation_root_path) + def _load(self): # By default, this prim cannot be instantiated from scratch! raise NotImplementedError("By default, an entity prim cannot be created from scratch.") @@ -763,6 +768,14 @@ def _denormalize_efforts(self, efforts, indices=None): """ return efforts * self.max_joint_efforts if indices is None else efforts * self.max_joint_efforts[indices] + def _verify_physics_view_is_valid(self): + """ + Helper function to make sure that the internal physics view is valid -- if not, will automatically refresh the + internal pointer + """ + if not self._physics_view.check(): + self._physics_view = og.sim.physics_sim_view.create_articulation_view(self.articulation_root_path) + def update_handles(self): """ Updates all internal handles for this prim, in case they change since initialization @@ -1278,6 +1291,59 @@ def aabb(self): return aabb_lo, aabb_hi + def get_coriolis_and_centrifugal_forces(self): + """ + 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!" + self._verify_physics_view_is_valid() + return self._physics_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!" + self._verify_physics_view_is_valid() + return self._physics_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!" + self._verify_physics_view_is_valid() + return self._physics_view.get_mass_matrices().reshape(self.n_dof, self.n_dof) + + def get_jacobian(self): + """ + 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!" + self._verify_physics_view_is_valid() + return self._physics_view.get_jacobians().squeeze(axis=0) + + def get_relative_jacobian(self): + """ + 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() + 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 + tf[:, 3:, 3:] = ori_t + return tf @ jac + def wake(self): """ Enable physics for this articulation diff --git a/omnigibson/prims/joint_prim.py b/omnigibson/prims/joint_prim.py index 8f8b2e0eb..0ab38b477 100644 --- a/omnigibson/prims/joint_prim.py +++ b/omnigibson/prims/joint_prim.py @@ -169,9 +169,11 @@ def _initialize(self): dof_props = self._dc.get_dof_properties(dof_handle) self._dof_handles.append(dof_handle) self._dof_properties.append(dof_props) - # Infer control type based on whether kp and kd are 0 or not + # Infer control type based on whether kp and kd are 0 or not, as well as whether this joint is driven or not kp, kd = dof_props.stiffness, dof_props.damping - if kp == 0.0: + if not self._driven: + control_type = ControlType.NONE + elif kp == 0.0: control_type = ControlType.EFFORT if kd == 0.0 else ControlType.VELOCITY else: control_type = ControlType.POSITION @@ -848,7 +850,6 @@ def set_effort(self, effort, normalized=False): # Sanity checks -- make sure that we're articulated (no control type check like position and velocity # because we can't set effort targets) and that we're driven self.assert_articulated() - assert self._driven, "Cannot set efforts for joint that is not driven!" # Standardize input effort = np.array([effort]) if self._n_dof == 1 and not isinstance(effort, Iterable) else np.array(effort) @@ -866,6 +867,9 @@ def keep_still(self): Zero out all velocities for this prim """ self.set_vel(np.zeros(self.n_dof)) + # If not driven, set torque equal to zero as well + if not self.driven: + self.set_effort(np.zeros(self.n_dof)) def _dump_state(self): pos, vel, effort = self.get_state() if self.articulated else (np.array([]), np.array([]), np.array([])) @@ -882,7 +886,7 @@ def _load_state(self, state): if self.articulated: self.set_pos(state["pos"], drive=False) self.set_vel(state["vel"], drive=False) - if self._driven: + if self.driven: self.set_effort(state["effort"]) if self._control_type == ControlType.POSITION: self.set_pos(state["target_pos"], drive=True) diff --git a/omnigibson/robots/manipulation_robot.py b/omnigibson/robots/manipulation_robot.py index 1468bbf0f..05b569558 100644 --- a/omnigibson/robots/manipulation_robot.py +++ b/omnigibson/robots/manipulation_robot.py @@ -150,6 +150,9 @@ def __init__( self._grasping_mode = grasping_mode self._disable_grasp_handling = disable_grasp_handling + # Other internal variables initialized later + self._eef_link_idxs = {arm: None for arm in self.arm_names} + # Initialize other variables used for assistive grasping self._ag_freeze_joint_pos = { arm: {} for arm in self.arm_names @@ -214,6 +217,13 @@ def _validate_configuration(self): def _initialize(self): super()._initialize() + + # Store index of EEF within link count + offset = -1 if self.fixed_base else 0 + link_path_to_idx = {path: i for i, path in enumerate(self._physics_view.link_paths[0])} + for arm in self.arm_names: + self._eef_link_idxs[arm] = link_path_to_idx[self.eef_links[arm].prim_path] + offset + if gm.AG_CLOTH: for arm in self.arm_names: self._ag_check_in_volume[arm], self._ag_calculate_volume[arm] = \ @@ -380,8 +390,15 @@ def get_control_dict(self): dic = super().get_control_dict() for arm in self.arm_names: - dic["eef_{}_pos_relative".format(arm)] = self.get_relative_eef_position(arm) - dic["eef_{}_quat_relative".format(arm)] = self.get_relative_eef_orientation(arm) + rel_eef_pos, rel_eef_quat = self.get_relative_eef_pose(arm) + dic[f"eef_{arm}_pos_relative"] = rel_eef_pos + dic[f"eef_{arm}_quat_relative"] = rel_eef_quat + dic[f"eef_{arm}_lin_vel_relative"] = self.get_relative_eef_lin_vel(arm) + dic[f"eef_{arm}_ang_vel_relative"] = self.get_relative_eef_ang_vel(arm) + # -n_joints because there may be an additional 6 entries at the beginning of the array, if this robot does + # not have a fixed base (i.e.: the 6DOF --> "floating" joint) + # see self.get_relative_jacobian() for more info + dic[f"eef_{arm}_jacobian_relative"] = self.get_relative_jacobian()[self._eef_link_idxs[arm], :, -self.n_joints:] return dic @@ -705,6 +722,33 @@ def get_relative_eef_orientation(self, arm="default"): arm = self.default_arm if arm == "default" else arm return self.get_relative_eef_pose(arm=arm)[1] + def get_relative_eef_lin_vel(self, arm="default"): + """ + Args: + arm (str): specific arm to grab relative eef linear velocity. + Default is "default" which corresponds to the first entry in self.arm_names + + + Returns: + 3-array: (x,y,z) Linear velocity of end-effector relative to robot base frame + """ + arm = self.default_arm if arm == "default" else arm + base_link_quat = self.get_orientation() + return T.quat2mat(base_link_quat).T @ self.eef_links[arm].get_linear_velocity() + + def get_relative_eef_ang_vel(self, arm="default"): + """ + Args: + arm (str): specific arm to grab relative eef angular velocity. + Default is "default" which corresponds to the first entry in self.arm_names + + Returns: + 3-array: (ax,ay,az) angular velocity of end-effector relative to robot base frame + """ + arm = self.default_arm if arm == "default" else arm + base_link_quat = self.get_orientation() + return T.quat2mat(base_link_quat).T @ self.eef_links[arm].get_angular_velocity() + def _calculate_in_hand_object_rigid(self, arm="default"): """ Calculates which object to assisted-grasp for arm @arm. Returns an (object_id, link_id) tuple or None @@ -861,6 +905,31 @@ def _default_arm_ik_controller_configs(self): } return dic + @property + def _default_arm_osc_controller_configs(self): + """ + Returns: + dict: Dictionary mapping arm appendage name to default controller config for an + operational space controller to control this robot's arm + """ + dic = {} + for arm in self.arm_names: + dic[arm] = { + "name": "OperationalSpaceController", + "task_name": f"eef_{arm}", + "control_freq": self._control_freq, + "default_joint_pos": self.default_joint_pos, + "control_limits": self.control_limits, + "dof_idx": self.arm_control_idx[arm], + "command_output_limits": ( + np.array([-0.2, -0.2, -0.2, -0.5, -0.5, -0.5]), + np.array([0.2, 0.2, 0.2, 0.5, 0.5, 0.5]), + ), + "mode": "pose_delta_ori", + "workspace_pose_limiter": None, + } + return dic + @property def _default_arm_null_joint_controller_configs(self): """ @@ -944,6 +1013,7 @@ def _default_controller_config(self): cfg = super()._default_controller_config arm_ik_configs = self._default_arm_ik_controller_configs + arm_osc_configs = self._default_arm_osc_controller_configs arm_joint_configs = self._default_arm_joint_controller_configs arm_null_joint_configs = self._default_arm_null_joint_controller_configs gripper_pj_configs = self._default_gripper_multi_finger_controller_configs @@ -954,6 +1024,7 @@ def _default_controller_config(self): for arm in self.arm_names: cfg["arm_{}".format(arm)] = { arm_ik_configs[arm]["name"]: arm_ik_configs[arm], + arm_osc_configs[arm]["name"]: arm_osc_configs[arm], arm_joint_configs[arm]["name"]: arm_joint_configs[arm], arm_null_joint_configs[arm]["name"]: arm_null_joint_configs[arm], } diff --git a/omnigibson/sensors/vision_sensor.py b/omnigibson/sensors/vision_sensor.py index 480f8c3a7..a0f620733 100644 --- a/omnigibson/sensors/vision_sensor.py +++ b/omnigibson/sensors/vision_sensor.py @@ -349,11 +349,29 @@ def clipping_range(self, limits): render() self.visible = True + @property + def horizontal_aperture(self): + """ + Returns: + float: horizontal aperture of this sensor, in mm + """ + return self.get_attribute("horizontalAperture") + + @horizontal_aperture.setter + def horizontal_aperture(self, length): + """ + Sets the focal length @length for this sensor + + Args: + length (float): horizontal aperture of this sensor, in meters + """ + self.set_attribute("horizontalAperture", length) + @property def focal_length(self): """ Returns: - float: focal length of this sensor, in meters + float: focal length of this sensor, in mm """ return self.get_attribute("focalLength") @@ -363,10 +381,33 @@ def focal_length(self, length): Sets the focal length @length for this sensor Args: - length (float): focal length of this sensor, in meters + length (float): focal length of this sensor, in mm """ self.set_attribute("focalLength", length) + @property + def intrinsic_matrix(self): + """ + Returns: + n-array: (3, 3) camera intrinsic matrix. Transforming point p (x,y,z) in the camera frame via K * p will + produce p' (x', y', w) - the point in the image plane. To get pixel coordiantes, divide x' and y' by w + """ + params = get_camera_params(viewport=self._viewport.viewport_api) + h, w = self.image_height, self.image_width + horizontal_fov = params["fov"] + vertical_fov = horizontal_fov * h / w + + f_x = (w / 2.0) / np.tan(horizontal_fov / 2.0) + f_y = (h / 2.0) / np.tan(vertical_fov / 2.0) + + K = np.array([ + [f_x, 0.0, w / 2.0], + [0.0, f_y, h / 2.0], + [0.0, 0.0, 1.0] + ]) + + return K + @property def _obs_space_mapping(self): # Generate the complex space types for special modalities: diff --git a/omnigibson/simulator.py b/omnigibson/simulator.py index 9d6f9e39f..d97320fad 100644 --- a/omnigibson/simulator.py +++ b/omnigibson/simulator.py @@ -1027,6 +1027,11 @@ def _open_new_stage(self): with suppress_omni_log(None): create_new_stage() + # Clear physics context + self._physics_context = None + if meets_minimum_isaac_version("2023.0.0"): + self._physx_fabric_interface = None + # Create world prim self.stage.DefinePrim("/World", "Xform") diff --git a/omnigibson/tasks/behavior_task.py b/omnigibson/tasks/behavior_task.py index ebfd8e3a1..083e895a4 100644 --- a/omnigibson/tasks/behavior_task.py +++ b/omnigibson/tasks/behavior_task.py @@ -18,6 +18,7 @@ from omnigibson.robots.robot_base import BaseRobot from omnigibson.systems.system_base import get_system, add_callback_on_system_init, add_callback_on_system_clear, \ REGISTERED_SYSTEMS +from omnigibson.scenes.scene_base import Scene from omnigibson.scenes.interactive_traversable_scene import InteractiveTraversableScene from omnigibson.utils.bddl_utils import OmniGibsonBDDLBackend, BDDLEntity, BEHAVIOR_ACTIVITIES, BDDLSampler from omnigibson.tasks.task_base import BaseTask @@ -143,7 +144,7 @@ def verify_scene_and_task_config(cls, scene_cfg, task_cfg): task_cfg.get("predefined_problem", None) is not None else task_cfg["activity_name"] if scene_file is None and scene_instance is None and not task_cfg["online_object_sampling"]: scene_instance = cls.get_cached_activity_scene_filename( - scene_model=scene_cfg["scene_model"], + scene_model=scene_cfg.get("scene_model", "Scene"), activity_name=activity_name, activity_definition_id=task_cfg.get("activity_definition_id", 0), activity_instance_id=task_cfg.get("activity_instance_id", 0), @@ -524,8 +525,8 @@ def name(self): @classproperty def valid_scene_types(cls): - # Must be an interactive traversable scene - return {InteractiveTraversableScene} + # Any scene can be used + return {Scene} @classproperty def default_termination_config(cls): diff --git a/omnigibson/tasks/task_base.py b/omnigibson/tasks/task_base.py index 6ec84aeb3..46d3eed04 100644 --- a/omnigibson/tasks/task_base.py +++ b/omnigibson/tasks/task_base.py @@ -50,6 +50,7 @@ def __init__(self, termination_config=None, reward_config=None): self._loaded = False self._reward = None self._done = None + self._success = None self._info = None self._low_dim_obs_dim = None @@ -160,7 +161,8 @@ def _reset_variables(self, env): """ # By default, reset reward, done, and info self._reward = None - self._done = None + self._done = False + self._success = False self._info = None def reset(self, env): @@ -311,6 +313,7 @@ def step(self, env, action): # Update the internal state of this task self._reward = reward self._done = done + self._success = done_info["success"] self._info = { "reward": reward_info, "done": done_info, @@ -344,6 +347,15 @@ def done(self): assert self._done is not None, "At least one step() must occur before done can be calculated!" return self._done + @property + def success(self): + """ + Returns: + bool: Whether this task has succeeded or not + """ + assert self._success is not None, "At least one step() must occur before success can be calculated!" + return self._success + @property def info(self): """ diff --git a/omnigibson/utils/bddl_utils.py b/omnigibson/utils/bddl_utils.py index c80f22b25..f7493e271 100644 --- a/omnigibson/utils/bddl_utils.py +++ b/omnigibson/utils/bddl_utils.py @@ -24,6 +24,7 @@ from omnigibson import object_states from omnigibson.object_states.factory import _KINEMATIC_STATE_SET from omnigibson.systems.system_base import is_system_active, get_system +from omnigibson.scenes.interactive_traversable_scene import InteractiveTraversableScene # Create module logger log = create_module_logger(module_name=__name__) @@ -276,7 +277,7 @@ def __init__( ): # Store internal variables from inputs self._env = env - self._scene_model = self._env.scene.scene_model + self._scene_model = self._env.scene.scene_model if isinstance(self._env.scene, InteractiveTraversableScene) else None self._agent = self._env.robots[0] if debug: gm.DEBUG = True @@ -411,9 +412,9 @@ def _parse_inroom_object_room_assignment(self): # Invalid room assignment return f"You have assigned room type for [{obj_synset}], but [{obj_synset}] is sampleable. " \ f"Only non-sampleable (scene) objects can have room assignment." - if room_type not in og.sim.scene.seg_map.room_sem_name_to_ins_name: + if self._scene_model is not None and room_type not in og.sim.scene.seg_map.room_sem_name_to_ins_name: # Missing room type - return f"Room type [{room_type}] missing in scene [{og.sim.scene.scene_model}]." + return f"Room type [{room_type}] missing in scene [{self._scene_model}]." if room_type not in self._room_type_to_object_instance: self._room_type_to_object_instance[room_type] = [] self._room_type_to_object_instance[room_type].append(obj_inst) @@ -521,6 +522,12 @@ def _build_sampling_order(self): # Sanity check kinematic objects -- any non-system must be kinematically sampled remaining_kinematic_entities = nonparticle_entities - unsampleable_obj_instances - \ self._inroom_object_instances - set.union(*(self._object_sampling_orders["kinematic"] + [set()])) + + # Possibly remove the agent entity if we're in an empty scene -- i.e.: no kinematic sampling needed for the + # agent + if self._scene_model is None: + remaining_kinematic_entities -= {"agent.n.01_1"} + if len(remaining_kinematic_entities) != 0: return f"Some objects do not have any kinematic condition defined for them in the initial conditions: " \ f"{', '.join(remaining_kinematic_entities)}" @@ -566,7 +573,8 @@ def _build_inroom_object_scope(self): valid_models = {cat: set(get_all_object_category_models_with_abilities(cat, abilities)) for cat in categories} - for room_inst in og.sim.scene.seg_map.room_sem_name_to_ins_name[room_type]: + room_insts = [None] if self._scene_model is None else og.sim.scene.seg_map.room_sem_name_to_ins_name[room_type] + for room_inst in room_insts: # A list of scene objects that satisfy the requested categories room_objs = og.sim.scene.object_registry("in_rooms", room_inst, default_val=[]) scene_objs = [obj for obj in room_objs if obj.category in categories and obj.model in valid_models[obj.category]] @@ -660,11 +668,15 @@ def _filter_object_scope(self, input_object_scope, conditions, condition_type): filtered_object_scope[room_type][scene_obj][room_inst].append(obj) # Compute most problematic objects - problematic_objs_by_proportion = defaultdict(list) - for child_scope_name, parent_obj_names in problematic_objs.items(): - problematic_objs_by_proportion[np.mean(list(parent_obj_names.values()))].append(child_scope_name) + if len(problematic_objs) == 0: + max_problematic_objs = [] + else: + problematic_objs_by_proportion = defaultdict(list) + for child_scope_name, parent_obj_names in problematic_objs.items(): + problematic_objs_by_proportion[np.mean(list(parent_obj_names.values()))].append(child_scope_name) + max_problematic_objs = problematic_objs_by_proportion[min(problematic_objs_by_proportion.keys())] - return filtered_object_scope, problematic_objs_by_proportion[min(problematic_objs_by_proportion.keys())] + return filtered_object_scope, max_problematic_objs def _consolidate_room_instance(self, filtered_object_scope, condition_type): """ diff --git a/omnigibson/utils/control_utils.py b/omnigibson/utils/control_utils.py index 8a327f1c8..64fa8398a 100644 --- a/omnigibson/utils/control_utils.py +++ b/omnigibson/utils/control_utils.py @@ -3,6 +3,7 @@ """ import lula import numpy as np +from numba import jit import omnigibson.utils.transform_utils as T from omnigibson.utils.sim_utils import meets_minimum_isaac_version @@ -129,4 +130,41 @@ def solve( if ik_results.success: return np.array(ik_results.cspace_position) else: - return None \ No newline at end of file + return None + + +@jit(nopython=True) +def orientation_error(desired, current): + """ + This function calculates a 3-dimensional orientation error vector for use in the + impedance controller. It does this by computing the delta rotation between the + inputs and converting that rotation to exponential coordinates (axis-angle + representation, where the 3d vector is axis * angle). + See https://en.wikipedia.org/wiki/Axis%E2%80%93angle_representation for more information. + Optimized function to determine orientation error from matrices + + Args: + desired (tensor): (..., 3, 3) where final two dims are 2d array representing target orientation matrix + current (tensor): (..., 3, 3) where final two dims are 2d array representing current orientation matrix + Returns: + tensor: (..., 3) where final dim is (ax, ay, az) axis-angle representing orientation error + """ + # convert input shapes + input_shape = desired.shape[:-2] + desired = desired.reshape(-1, 3, 3) + current = current.reshape(-1, 3, 3) + + # grab relevant info + rc1 = current[:, :, 0] + rc2 = current[:, :, 1] + rc3 = current[:, :, 2] + rd1 = desired[:, :, 0] + rd2 = desired[:, :, 1] + rd3 = desired[:, :, 2] + + error = 0.5 * (np.cross(rc1, rd1) + np.cross(rc2, rd2) + np.cross(rc3, rd3)) + + # Reshape + error = error.reshape(*input_shape, 3) + + return error diff --git a/omnigibson/utils/python_utils.py b/omnigibson/utils/python_utils.py index 7fcdf146c..1b3c63e3e 100644 --- a/omnigibson/utils/python_utils.py +++ b/omnigibson/utils/python_utils.py @@ -824,6 +824,26 @@ def __setattr__(self, key, value): super().__setattr__(key, value) +def nums2array(nums, dim, dtype=float): + """ + Converts input @nums into numpy array of length @dim. If @nums is a single number, broadcasts input to + corresponding dimension size @dim before converting into numpy array + + Args: + nums (float or array): Numbers to map to numpy array + dim (int): Size of array to broadcast input to + + Returns: + torch.Tensor: Mapped input numbers + """ + # Make sure the inputted nums isn't a string + assert not isinstance(nums, str), "Only numeric types are supported for this operation!" + + out = np.array(nums, dtype=dtype) if isinstance(nums, Iterable) else np.ones(dim, dtype=dtype) * nums + + return out + + def clear(): """ Clear state tied to singleton classes diff --git a/omnigibson/utils/ui_utils.py b/omnigibson/utils/ui_utils.py index b0fea3392..4cfc8518b 100644 --- a/omnigibson/utils/ui_utils.py +++ b/omnigibson/utils/ui_utils.py @@ -596,6 +596,34 @@ def generate_ik_keypress_mapping(self, controller_info): return mapping + def generate_osc_keypress_mapping(self, controller_info): + """ + Generates a dictionary for keypress mappings for OSC control, based on the inputted @controller_info + + Args: + controller_info (dict): Dictionary of controller information for the specific robot arm to control + with OSC + + Returns: + dict: Populated keypress mappings for IK to control the specified controller + """ + mapping = {} + + mapping[carb.input.KeyboardInput.UP] = {"idx": controller_info["start_idx"] + 0, "val": 0.5} + mapping[carb.input.KeyboardInput.DOWN] = {"idx": controller_info["start_idx"] + 0, "val": -0.5} + mapping[carb.input.KeyboardInput.RIGHT] = {"idx": controller_info["start_idx"] + 1, "val": -0.5} + mapping[carb.input.KeyboardInput.LEFT] = {"idx": controller_info["start_idx"] + 1, "val": 0.5} + mapping[carb.input.KeyboardInput.P] = {"idx": controller_info["start_idx"] + 2, "val": 0.5} + mapping[carb.input.KeyboardInput.SEMICOLON] = {"idx": controller_info["start_idx"] + 2, "val": -0.5} + mapping[carb.input.KeyboardInput.N] = {"idx": controller_info["start_idx"] + 3, "val": 0.5} + mapping[carb.input.KeyboardInput.B] = {"idx": controller_info["start_idx"] + 3, "val": -0.5} + mapping[carb.input.KeyboardInput.O] = {"idx": controller_info["start_idx"] + 4, "val": 0.5} + mapping[carb.input.KeyboardInput.U] = {"idx": controller_info["start_idx"] + 4, "val": -0.5} + mapping[carb.input.KeyboardInput.V] = {"idx": controller_info["start_idx"] + 5, "val": 0.5} + mapping[carb.input.KeyboardInput.C] = {"idx": controller_info["start_idx"] + 5, "val": -0.5} + + return mapping + def populate_keypress_mapping(self): """ Populates the mapping @self.keypress_mapping, which maps keypresses to action info: @@ -629,6 +657,9 @@ def populate_keypress_mapping(self): elif info["name"] == "InverseKinematicsController": self.ik_arms.append(component) self.keypress_mapping.update(self.generate_ik_keypress_mapping(controller_info=info)) + elif info["name"] == "OperationalSpaceController": + self.ik_arms.append(component) + self.keypress_mapping.update(self.generate_osc_keypress_mapping(controller_info=info)) elif info["name"] == "MultiFingerGripperController": if info["command_dim"] > 1: for i in range(info["command_dim"]): diff --git a/omnigibson/utils/usd_utils.py b/omnigibson/utils/usd_utils.py index d20512c14..f8fe10c8a 100644 --- a/omnigibson/utils/usd_utils.py +++ b/omnigibson/utils/usd_utils.py @@ -124,7 +124,7 @@ def get_camera_params(viewport): view_proj_mat = helpers.get_view_proj_mat(view_params) return { - "pose": np.array(prim_tf), + "pose": np.array(prim_tf).T, # omni natively gives transposed pose so we have to "un"-transpose it "fov": fov, "focal_length": view_params["focal_length"], "horizontal_aperture": view_params["horizontal_aperture"],