import_custom_robot

Helper script to download BEHAVIOR-1K dataset and assets.

add_sensor(stage, root_prim, sensor_type, link_name, parent_link_name=None, pos_offset=None, ori_offset=None)

Adds sensor to robot. This is an in-place operation on @root_prim

Parameters:

Name	Type	Description	Default
stage	Stage	Current active omniverse stage	required
root_prim	Prim	Root prim of the current robot, assumed to be on the current stage	required
sensor_type	str	Sensor to create. Valid options are: {Camera, Lidar, VisualSphere}	required
link_name	str	Link to attach the created sensor prim to. If this link does not already exist in the robot's current set of links, a new one will be created as a child of @parent_link_name's link	required
parent_link_name	None or str	If specified, parent link from which to create a new child link @link_name. If set, @link_name should NOT be a link already found on the robot!	None
pos_offset	None or 3 - tuple	If specified, (x,y,z) local translation offset to apply. If @parent_link_name is specified, defines offset of @link_name wrt @parent_link_name If only @link_name is specified, defines offset of the sensor prim wrt @link_name	None
ori_offset	None or 3 - tuple	If specified, (x,y,z,w) quaternion rotation offset to apply. If @parent_link_name is specified, defines offset of @link_name wrt @parent_link_name If only @link_name is specified, defines offset of the sensor prim wrt @link_name	None

Source code in OmniGibson/omnigibson/examples/robots/import_custom_robot.py

def add_sensor(stage, root_prim, sensor_type, link_name, parent_link_name=None, pos_offset=None, ori_offset=None):
    """
    Adds sensor to robot. This is an in-place operation on @root_prim

    Args:
        stage (Usd.Stage): Current active omniverse stage
        root_prim (Usd.Prim): Root prim of the current robot, assumed to be on the current stage
        sensor_type (str): Sensor to create. Valid options are: {Camera, Lidar, VisualSphere}
        link_name (str): Link to attach the created sensor prim to. If this link does not already exist in the robot's
            current set of links, a new one will be created as a child of @parent_link_name's link
        parent_link_name (None or str): If specified, parent link from which to create a new child link @link_name. If
            set, @link_name should NOT be a link already found on the robot!
        pos_offset (None or 3-tuple): If specified, (x,y,z) local translation offset to apply.
            If @parent_link_name is specified, defines offset of @link_name wrt @parent_link_name
            If only @link_name is specified, defines offset of the sensor prim wrt @link_name
        ori_offset (None or 3-tuple): If specified, (x,y,z,w) quaternion rotation offset to apply.
            If @parent_link_name is specified, defines offset of @link_name wrt @parent_link_name
            If only @link_name is specified, defines offset of the sensor prim wrt @link_name
    """
    # Make sure pos and ori offsets are defined
    if pos_offset is None or pos_offset == {}:  # May be {} from empty addict key
        pos_offset = (0, 0, 0)
    if ori_offset is None or ori_offset == {}:  # May be {} from empty addict key
        ori_offset = (0, 0, 0, 1)

    pos_offset = th.tensor(pos_offset, dtype=th.float)
    ori_offset = th.tensor(ori_offset, dtype=th.float)

    # Sanity check link / parent link combos
    root_prim_path = root_prim.GetPrimPath().pathString
    if parent_link_name is None or parent_link_name == {}:  # May be {} from empty addict key
        parent_link_prim = None
    else:
        parent_path = f"{root_prim_path}/{parent_link_name}"
        assert lazy.isaacsim.core.utils.prims.is_prim_path_valid(
            parent_path
        ), f"Could not find parent link within robot with name {parent_link_name}!"
        parent_link_prim = lazy.isaacsim.core.utils.prims.get_prim_at_path(parent_path)

    # If parent link is defined, link prim should NOT exist (this should be a new link)
    link_prim_path = f"{root_prim_path}/{link_name}"
    link_prim_exists = lazy.isaacsim.core.utils.prims.is_prim_path_valid(link_prim_path)
    if parent_link_prim is not None:
        assert not link_prim_exists, f"Since parent link is defined, link_name {link_name} must be a link that is NOT pre-existing within the robot's set of links!"
        # Manually create a new prim (specified offset)
        create_rigid_prim(
            stage=stage,
            link_prim_path=link_prim_path,
        )
        link_prim_xform = lazy.isaacsim.core.prims.xform_prim.XFormPrim(prim_path=link_prim_path)

        # Create fixed joint to connect the two together
        create_joint(
            prim_path=f"{parent_path}/{parent_link_name}_{link_name}_joint",
            joint_type="FixedJoint",
            body0=parent_path,
            body1=link_prim_path,
            joint_frame_in_parent_frame_pos=pos_offset,
            joint_frame_in_parent_frame_quat=ori_offset,
        )

        # Set child prim to the appropriate local pose -- this should be the parent local pose transformed by
        # the additional offset
        parent_prim_xform = lazy.isaacsim.core.prims.xform_prim.XFormPrim(prim_path=parent_path)
        parent_pos, parent_quat = parent_prim_xform.get_local_pose()
        parent_quat = parent_quat[[1, 2, 3, 0]]
        parent_pose = T.pose2mat((th.tensor(parent_pos), th.tensor(parent_quat)))
        offset_pose = T.pose2mat((pos_offset, ori_offset))
        child_pose = parent_pose @ offset_pose
        link_pos, link_quat = T.mat2pose(child_pose)
        link_prim_xform.set_local_pose(link_pos, link_quat[[3, 0, 1, 2]])

    else:
        # Otherwise, link prim MUST exist
        assert link_prim_exists, f"Since no parent link is defined, link_name {link_name} must be a link that IS pre-existing within the robot's set of links!"

    # Define functions to generate the desired sensor prim
    if sensor_type == "Camera":
        create_sensor_prim = lambda parent_prim_path: lazy.pxr.UsdGeom.Camera.Define(
            stage, f"{parent_prim_path}/Camera"
        ).GetPrim()
    elif sensor_type == "Lidar":
        create_sensor_prim = lambda parent_prim_path: lazy.omni.kit.commands.execute(
            "RangeSensorCreateLidar",
            path="/Lidar",
            parent=parent_prim_path,
            min_range=0.4,
            max_range=100.0,
            draw_points=False,
            draw_lines=False,
            horizontal_fov=360.0,
            vertical_fov=30.0,
            horizontal_resolution=0.4,
            vertical_resolution=4.0,
            rotation_rate=0.0,
            high_lod=False,
            yaw_offset=0.0,
            enable_semantics=False,
        )[1].GetPrim()
    elif sensor_type == "VisualSphere":
        create_sensor_prim = lambda parent_prim_path: create_primitive_mesh(
            prim_path=f"{parent_prim_path}/VisualSphere",
            primitive_type="Sphere",
            extents=0.01,
            stage=stage,
        ).GetPrim()
    else:
        raise ValueError(f"Got unknown sensor type: {sensor_type}!")

    # Create the new prim as a child of the link prim
    sensor_prim = create_sensor_prim(parent_prim_path=link_prim_path)
    _add_xform_properties(sensor_prim)

    # If sensor prim is a camera, set some default values
    if sensor_type == "Camera":
        sensor_prim.GetAttribute("focalLength").Set(17.0)
        sensor_prim.GetAttribute("clippingRange").Set(lazy.pxr.Gf.Vec2f(0.001, 1000000.0))
        # Refresh visibility
        lazy.pxr.UsdGeom.Imageable(sensor_prim).MakeInvisible()
        og.sim.render()
        lazy.pxr.UsdGeom.Imageable(sensor_prim).MakeVisible()

    # If we didn't have a parent prim defined, we need to add the offset directly to this sensor
    if parent_link_prim is None:
        sensor_prim.GetAttribute("xformOp:translate").Set(lazy.pxr.Gf.Vec3d(*pos_offset.tolist()))
        sensor_prim.GetAttribute("xformOp:orient").Set(
            lazy.pxr.Gf.Quatd(*ori_offset[[3, 0, 1, 2]].tolist())
        )  # expects (w,x,y,z)

create_curobo_cfgs(robot_prim, robot_urdf_path, curobo_cfg, root_link, save_dir, is_holonomic=False)

Creates a set of curobo configs based on @robot_prim and @curobo_cfg

Parameters:

Name	Type	Description	Default
robot_prim	Prim	Top-level prim defining the robot in the current USD stage	required
robot_urdf_path	str	Path to robot URDF file	required
curobo_cfg	Dict	Dictionary of relevant curobo information	required
root_link	str	Name of the robot's root link, BEFORE any holonomic joints are applied	required
save_dir	str	Path to the directory to save generated curobo files	required
is_holonomic	bool	Whether the robot has a holonomic base applied or not	False

Source code in OmniGibson/omnigibson/examples/robots/import_custom_robot.py

def create_curobo_cfgs(robot_prim, robot_urdf_path, curobo_cfg, root_link, save_dir, is_holonomic=False):
    """
    Creates a set of curobo configs based on @robot_prim and @curobo_cfg

    Args:
        robot_prim (Usd.Prim): Top-level prim defining the robot in the current USD stage
        robot_urdf_path (str): Path to robot URDF file
        curobo_cfg (Dict): Dictionary of relevant curobo information
        root_link (str): Name of the robot's root link, BEFORE any holonomic joints are applied
        save_dir (str): Path to the directory to save generated curobo files
        is_holonomic (bool): Whether the robot has a holonomic base applied or not
    """
    robot_name = robot_prim.GetName()

    # Left, then right by default if sorted alphabetically
    ee_links = list(sorted(curobo_cfg.eef_to_gripper_info.keys()))

    # Find all joints that have a negative axis specified so we know to flip them in curobo
    tree = ET.parse(robot_urdf_path)
    root = tree.getroot()
    flip_joints = dict()
    flip_joint_limits = []
    for joint in root.findall("joint"):
        if joint.attrib["type"] != "fixed":
            axis = th.round(_space_string_to_tensor(joint.find("axis").attrib["xyz"]))
            axis_idx = th.nonzero(axis).squeeze().item()
            flip_joints[joint.attrib["name"]] = "XYZ"[axis_idx]
            is_negative = (axis[axis_idx] < 0).item()
            if is_negative:
                flip_joint_limits.append(joint.attrib["name"])

    def get_joint_upper_limit(root_prim, joint_name):
        joint_prim = find_prim_with_name(name=joint_name, root_prim=root_prim)
        assert joint_prim is not None, f"Could not find joint prim with name {joint_name}!"
        return joint_prim.GetAttribute("physics:upperLimit").Get()

    # The original format from Lula is a list of dicts, so we need to convert to a single dict
    if isinstance(curobo_cfg.collision_spheres, list):
        collision_spheres = {k: v for c in curobo_cfg.collision_spheres for k, v in c.to_dict().items()}
    else:
        collision_spheres = curobo_cfg.collision_spheres.to_dict()

    # Generate list of collision link names -- this is simply the list of all link names from the
    # collision spheres specification
    all_collision_link_names = list(collision_spheres.keys())

    joint_prims = find_all_joints(robot_prim, only_articulated=True)
    all_joint_names = [joint_prim.GetName() for joint_prim in joint_prims]
    lock_joints = curobo_cfg.lock_joints.to_dict() if curobo_cfg.lock_joints else {}
    if is_holonomic:
        # Move the final six joints to the beginning, since the holonomic joints are added at the end
        all_joint_names = list(reversed(all_joint_names[-6:])) + all_joint_names[:-6]
        lock_joints["base_footprint_z_joint"] = None
        lock_joints["base_footprint_rx_joint"] = None
        lock_joints["base_footprint_ry_joint"] = None

    default_generated_cfg = {
        "usd_robot_root": f"/{robot_prim.GetName()}",
        "usd_flip_joints": flip_joints,
        "usd_flip_joint_limits": flip_joint_limits,
        "base_link": "base_footprint_x" if is_holonomic else root_link,
        "ee_link": ee_links[0],
        "link_names": ee_links[1:],
        "lock_joints": lock_joints,
        "extra_links": {},
        "collision_link_names": deepcopy(all_collision_link_names),
        "collision_spheres": collision_spheres,
        "collision_sphere_buffer": 0.002,
        "extra_collision_spheres": {},
        "self_collision_ignore": curobo_cfg.self_collision_ignore.to_dict(),
        "self_collision_buffer": curobo_cfg.self_collision_buffer.to_dict(),
        "use_global_cumul": True,
        "mesh_link_names": deepcopy(all_collision_link_names),
        "external_asset_path": None,
        "cspace": {
            "joint_names": all_joint_names,
            "retract_config": None,
            "null_space_weight": [1] * len(all_joint_names),
            "cspace_distance_weight": [1] * len(all_joint_names),
            "max_jerk": 500.0,
            "max_acceleration": 15.0,
        },
    }

    for eef_link_name, gripper_info in curobo_cfg.eef_to_gripper_info.items():
        attached_obj_link_name = f"attached_object_{eef_link_name}"
        for jnt_name in gripper_info["joints"]:
            default_generated_cfg["lock_joints"][jnt_name] = None
        default_generated_cfg["extra_links"][attached_obj_link_name] = {
            "parent_link_name": eef_link_name,
            "link_name": attached_obj_link_name,
            "fixed_transform": [0, 0, 0, 1, 0, 0, 0],  # (x,y,z,w,x,y,z)
            "joint_type": "FIXED",
            "joint_name": f"{attached_obj_link_name}_joint",
        }
        default_generated_cfg["collision_link_names"].append(attached_obj_link_name)
        default_generated_cfg["extra_collision_spheres"][attached_obj_link_name] = 32
        for link_name in gripper_info["links"]:
            if link_name not in default_generated_cfg["self_collision_ignore"]:
                default_generated_cfg["self_collision_ignore"][link_name] = []
            default_generated_cfg["self_collision_ignore"][link_name].append(attached_obj_link_name)
        default_generated_cfg["mesh_link_names"].append(attached_obj_link_name)

    # Save generated file
    Path(save_dir).mkdir(parents=True, exist_ok=True)
    save_fpath = f"{save_dir}/{robot_name}_description_curobo_default.yaml"
    with open(save_fpath, "w+") as f:
        yaml.dump({"robot_cfg": {"kinematics": default_generated_cfg}}, f)

    # Permute the default config to have additional base only / arm only configs
    # Only relevant if robot is holonomic
    if is_holonomic:
        # Create base only config
        base_only_cfg = deepcopy(default_generated_cfg)
        base_only_cfg["ee_link"] = root_link
        base_only_cfg["link_names"] = []
        for jnt_name in all_joint_names:
            if jnt_name not in base_only_cfg["lock_joints"] and "base_footprint" not in jnt_name:
                # Lock this joint
                base_only_cfg["lock_joints"][jnt_name] = None
        save_base_fpath = f"{save_dir}/{robot_name}_description_curobo_base.yaml"
        with open(save_base_fpath, "w+") as f:
            yaml.dump({"robot_cfg": {"kinematics": base_only_cfg}}, f)

        # Create arm only config
        arm_only_cfg = deepcopy(default_generated_cfg)
        for jnt_name in {"base_footprint_x_joint", "base_footprint_y_joint", "base_footprint_rz_joint"}:
            arm_only_cfg["lock_joints"][jnt_name] = None
        save_arm_fpath = f"{save_dir}/{robot_name}_description_curobo_arm.yaml"
        with open(save_arm_fpath, "w+") as f:
            yaml.dump({"robot_cfg": {"kinematics": arm_only_cfg}}, f)

        # Create arm only no torso config
        arm_only_no_torso_cfg = deepcopy(arm_only_cfg)
        for jnt_name in curobo_cfg.torso_joints:
            arm_only_no_torso_cfg["lock_joints"][jnt_name] = None
        save_arm_no_torso_fpath = f"{save_dir}/{robot_name}_description_curobo_arm_no_torso.yaml"
        with open(save_arm_no_torso_fpath, "w+") as f:
            yaml.dump({"robot_cfg": {"kinematics": arm_only_no_torso_cfg}}, f)

create_rigid_prim(stage, link_prim_path)

Creates a new rigid link prim nested under @root_prim

Parameters:

Name	Type	Description	Default
stage	Stage	Current active omniverse stage	required
link_prim_path	str	Prim path at which link will be created. Should not already exist on the stage	required

Returns:

Type	Description
Prim	Newly created rigid prim

Source code in OmniGibson/omnigibson/examples/robots/import_custom_robot.py

def create_rigid_prim(stage, link_prim_path):
    """
    Creates a new rigid link prim nested under @root_prim

    Args:
        stage (Usd.Stage): Current active omniverse stage
        link_prim_path (str): Prim path at which link will be created. Should not already exist on the stage

    Returns:
        Usd.Prim: Newly created rigid prim
    """
    # Make sure link prim does NOT already exist (this should be a new link)
    link_prim_exists = stage.GetPrimAtPath(link_prim_path).IsValid()
    assert (
        not link_prim_exists
    ), f"Cannot create new link because there already exists a link at prim path {link_prim_path}!"

    # Manually create a new prim (specified offset)
    link_prim = lazy.pxr.UsdGeom.Xform.Define(stage, link_prim_path).GetPrim()
    _add_xform_properties(prim=link_prim)

    # Add rigid prim API to new link prim
    lazy.pxr.UsdPhysics.RigidBodyAPI.Apply(link_prim)
    lazy.pxr.PhysxSchema.PhysxRigidBodyAPI.Apply(link_prim)

    return link_prim

find_articulation_root_prim(root_prim)

Recursively searches children of @root_prim to find the articulation root

Parameters:

Name	Type	Description	Default
root_prim	Prim	Root prim to search	required

Returns:

Type	Description
None or Prim	If found, articulation root prim

Source code in OmniGibson/omnigibson/examples/robots/import_custom_robot.py

def find_articulation_root_prim(root_prim):
    """
    Recursively searches children of @root_prim to find the articulation root

    Args:
        root_prim (Usd.Prim): Root prim to search

    Returns:
        None or Usd.Prim: If found, articulation root prim
    """
    return _find_prim_with_condition(
        condition=lambda prim: prim.HasAPI(lazy.pxr.UsdPhysics.ArticulationRootAPI), root_prim=root_prim
    )

find_prim_with_name(name, root_prim)

Recursively searches children of @root_prim to find first instance of prim including string @name

Parameters:

Name	Type	Description	Default
name	str	Name of the prim to search	required
root_prim	Prim	Root prim to search	required

Returns:

Type	Description
None or Prim	If found, first prim whose prim name includes @name

Source code in OmniGibson/omnigibson/examples/robots/import_custom_robot.py

def find_prim_with_name(name, root_prim):
    """
    Recursively searches children of @root_prim to find first instance of prim including string @name

    Args:
        name (str): Name of the prim to search
        root_prim (Usd.Prim): Root prim to search

    Returns:
        None or Usd.Prim: If found, first prim whose prim name includes @name
    """
    return _find_prim_with_condition(condition=lambda prim: name in prim.GetName(), root_prim=root_prim)

make_joint_fixed(stage, root_prim, joint_name)

Converts a revolute / prismatic joint @joint_name into a fixed joint

NOTE: This is an in-place operation!

Parameters:

Name	Type	Description	Default
stage	Stage	Current active omniverse stage	required
root_prim	Prim	Root prim of the current robot, assumed to be on the current stage	required
joint_name	str	Joint to convert to be fixed	required

Source code in OmniGibson/omnigibson/examples/robots/import_custom_robot.py

def make_joint_fixed(stage, root_prim, joint_name):
    """
    Converts a revolute / prismatic joint @joint_name into a fixed joint

    NOTE: This is an in-place operation!

    Args:
        stage (Usd.Stage): Current active omniverse stage
        root_prim (Usd.Prim): Root prim of the current robot, assumed to be on the current stage
        joint_name (str): Joint to convert to be fixed
    """
    joint_prim = find_prim_with_name(name=joint_name, root_prim=root_prim)
    assert joint_prim is not None, f"Could not find joint prim with name {joint_name}!"

    # Remove its Joint APIs and add Fixed Joint API
    joint_type = joint_prim.GetTypeName()
    if joint_type != "PhysicsFixedJoint":
        assert joint_type in {
            "PhysicsRevoluteJoint",
            "PhysicsPrismaticJoint",
        }, f"Got invalid joint type: {joint_type}. Only PhysicsRevoluteJoint and PhysicsPrismaticJoint are supported!"

        lazy.omni.kit.commands.execute("RemovePhysicsComponentCommand", usd_prim=joint_prim, component=joint_type)
        lazy.pxr.UsdPhysics.FixedJoint.Define(stage, joint_prim.GetPrimPath().pathString)

set_link_collision_approximation(stage, root_prim, link_name, approximation_type)

Sets all collision geoms under @link_name to be @approximation type Args: approximation_type (str): approximation used for collision. Can be one of: {"none", "convexHull", "convexDecomposition", "meshSimplification", "sdf", "boundingSphere", "boundingCube"} If None, the approximation will use the underlying triangle mesh.

Source code in OmniGibson/omnigibson/examples/robots/import_custom_robot.py

def set_link_collision_approximation(stage, root_prim, link_name, approximation_type):
    """
    Sets all collision geoms under @link_name to be @approximation type
    Args:
        approximation_type (str): approximation used for collision.
            Can be one of: {"none", "convexHull", "convexDecomposition", "meshSimplification", "sdf",
                "boundingSphere", "boundingCube"}
            If None, the approximation will use the underlying triangle mesh.
    """
    # Sanity check approximation type
    assert_valid_key(
        key=approximation_type,
        valid_keys={
            "none",
            "convexHull",
            "convexDecomposition",
            "meshSimplification",
            "sdf",
            "boundingSphere",
            "boundingCube",
        },
        name="collision approximation type",
    )

    # Find the link prim first
    link_prim = find_prim_with_name(name=link_name, root_prim=root_prim)
    assert link_prim is not None, f"Could not find link prim with name {link_name}!"

    # Iterate through all children that are mesh prims
    mesh_prims = find_all_prim_children_with_type(prim_type="Mesh", root_prim=link_prim)

    # For each mesh prim, check if it is collision -- if so, update the approximation type appropriately
    for mesh_prim in mesh_prims:
        if not mesh_prim.HasAPI(lazy.pxr.UsdPhysics.MeshCollisionAPI):
            # This is a visual mesh, so skip
            continue
        mesh_collision_api = lazy.pxr.UsdPhysics.MeshCollisionAPI(mesh_prim)

        # Make sure to add the appropriate API if we're setting certain values
        if approximation_type == "convexHull" and not mesh_prim.HasAPI(
            lazy.pxr.PhysxSchema.PhysxConvexHullCollisionAPI
        ):
            lazy.pxr.PhysxSchema.PhysxConvexHullCollisionAPI.Apply(mesh_prim)
        elif approximation_type == "convexDecomposition" and not mesh_prim.HasAPI(
            lazy.pxr.PhysxSchema.PhysxConvexDecompositionCollisionAPI
        ):
            lazy.pxr.PhysxSchema.PhysxConvexDecompositionCollisionAPI.Apply(mesh_prim)
        elif approximation_type == "meshSimplification" and not mesh_prim.HasAPI(
            lazy.pxr.PhysxSchema.PhysxTriangleMeshSimplificationCollisionAPI
        ):
            lazy.pxr.PhysxSchema.PhysxTriangleMeshSimplificationCollisionAPI.Apply(mesh_prim)
        elif approximation_type == "sdf" and not mesh_prim.HasAPI(lazy.pxr.PhysxSchema.PhysxSDFMeshCollisionAPI):
            lazy.pxr.PhysxSchema.PhysxSDFMeshCollisionAPI.Apply(mesh_prim)
        elif approximation_type == "none" and not mesh_prim.HasAPI(lazy.pxr.PhysxSchema.PhysxTriangleMeshCollisionAPI):
            lazy.pxr.PhysxSchema.PhysxTriangleMeshCollisionAPI.Apply(mesh_prim)

        if approximation_type == "convexHull":
            pch_api = lazy.pxr.PhysxSchema.PhysxConvexHullCollisionAPI(mesh_prim)
            # Also make sure the maximum vertex count is 60 (max number compatible with GPU)
            # https://docs.omniverse.nvidia.com/app_create/prod_extensions/ext_physics/rigid-bodies.html#collision-settings
            if pch_api.GetHullVertexLimitAttr().Get() is None:
                pch_api.CreateHullVertexLimitAttr()
            pch_api.GetHullVertexLimitAttr().Set(60)

        mesh_collision_api.GetApproximationAttr().Set(approximation_type)