init commit

workflows/simbox/core/objects/__init__.py

@@ -0,0 +1,34 @@
+"""Objects module initialization."""
+
+from core.objects.base_object import OBJECT_DICT
+
+from .articulated_object import ArticulatedObject
+from .conveyor_object import ConveyorObject
+from .geometry_object import GeometryObject
+from .plane_object import PlaneObject
+from .rigid_object import RigidObject
+from .shape_object import ShapeObject
+from .xform_object import XFormObject
+
+# Explicitly declare the public interface
+__all__ = [
+    "ArticulatedObject",
+    "ConveyorObject",
+    "GeometryObject",
+    "PlaneObject",
+    "RigidObject",
+    "ShapeObject",
+    "XFormObject",
+    "get_object_cls",
+    "get_object_dict",
+]
+
+
+def get_object_cls(category_name):
+    """Get object class by category name."""
+    return OBJECT_DICT[category_name]
+
+
+def get_object_dict():
+    """Get object dictionary."""
+    return OBJECT_DICT

workflows/simbox/core/objects/articulated_object.py

@@ -0,0 +1,157 @@
+import glob
+import json
+import os
+import random
+
+import numpy as np
+from core.objects.base_object import register_object
+from omni.isaac.core.articulations.articulation import Articulation
+from omni.isaac.core.utils.stage import add_reference_to_stage
+
+try:
+    from omni.isaac.core.materials.omni_pbr import OmniPBR  # Isaac Sim 4.1.0 / 4.2.0
+except ImportError:
+    from isaacsim.core.api.materials import OmniPBR  # Isaac Sim 4.5.0
+
+
+@register_object
+class ArticulatedObject(Articulation):
+    def __init__(self, asset_root, root_prim_path, cfg, *args, **kwargs):
+        self.asset_root = asset_root
+        self.object_name = cfg["name"]
+        self.asset_relative_path = cfg["path"]
+        self.object_dir = os.path.join(asset_root, cfg["path"])
+        self._root_prim_path = root_prim_path
+        object_info_path = os.path.join(asset_root, cfg["obj_info_path"])
+        with open(object_info_path, "r", encoding="utf-8") as f:
+            object_info = json.load(f)
+        self.category = cfg["category"]
+        self.cfg = cfg
+        self.get_articulated_info(object_info)
+        super().__init__(prim_path=self.object_prim_path, name=cfg["name"], *args, **kwargs)
+
+    def update_articulated_info(self, obj_info_path):
+        object_info_path = f"{self.asset_root}/{obj_info_path}"
+        with open(object_info_path, "r", encoding="utf-8") as f:
+            object_info = json.load(f)
+
+        prim_path_list = ["object_prim_path", "object_link_path", "object_base_path"]
+        axis_list = [
+            "object_link0_rot_axis",
+            "object_link0_move_axis",
+            "object_link0_contact_axis",
+            "object_base_front_axis",
+        ]
+        joint_path_list = ["object_joint_path", "object_revolute_joint_path"]
+        joint_idx_list = ["joint_index", "object_revolute_joint_idx"]
+
+        self.object_keypoints = object_info["object_keypoints"]  # Object keypoints in link_0 frame
+        for key, item in self.object_keypoints.items():
+            self.object_keypoints[key] = np.append(item, [1.0], axis=0)
+
+        self.object_scale = np.array(object_info["object_scale"])
+
+        for key, item in object_info.items():
+            if key in prim_path_list:
+                setattr(self, key, self._root_prim_path + object_info[key])
+            elif key in axis_list:
+                setattr(self, key, object_info[key])
+            elif key in joint_path_list:
+                self.object_joint_path = self._root_prim_path + object_info[key]
+            elif key in joint_idx_list:
+                self.object_joint_index = object_info[key]
+
+        self.articulation_initial_joint_position = self._articulation_view.get_joint_positions()[
+            :, self.object_joint_index
+        ]
+
+        self.contact_plane_normal = None
+
+    def get_articulated_info(self, object_info):
+        prim_path_list = ["object_prim_path", "object_link_path", "object_base_path"]
+        axis_list = [
+            "object_link0_rot_axis",
+            "object_link0_move_axis",
+            "object_link0_contact_axis",
+            "object_base_front_axis",
+        ]
+        joint_path_list = ["object_joint_path", "object_revolute_joint_path"]
+        joint_idx_list = ["joint_index", "object_revolute_joint_idx"]
+        self.object_keypoints = object_info["object_keypoints"]  # Object keypoints in link_0 frame
+        for key, item in self.object_keypoints.items():
+            self.object_keypoints[key] = np.append(item, [1.0], axis=0)
+        self.object_scale = np.array(object_info["object_scale"])
+        self.object_usd = os.path.join(self.object_dir, "instance.usd")
+        for key, item in object_info.items():
+            if key in prim_path_list:
+                setattr(self, key, self._root_prim_path + object_info[key])
+            elif key in axis_list:
+                setattr(self, key, object_info[key])
+            elif key in joint_path_list:
+                self.object_joint_path = self._root_prim_path + object_info[key]
+            elif key in joint_idx_list:
+                self.object_joint_index = object_info[key]
+
+        # Compute joint number
+        self.object_joint_number = 0
+        # Contact plane normal
+        self.contact_plane_normal = None
+        add_reference_to_stage(usd_path=self.object_usd, prim_path=self.object_prim_path)
+
+    def get_joint_position(self, stage):
+        joint_parent_prim = stage.GetPrimAtPath(self.object_joint_path.rsplit("/", 1)[0])
+        for child in joint_parent_prim.GetAllChildren():
+            if child.GetTypeName() == "PhysicsPrismaticJoint" or child.GetTypeName() == "PhysicsRevoluteJoint":
+                self.object_joint_number += 1
+
+        # Fix the asset base
+        if self.cfg.get("fix_base", False):
+            parent_prim_path = os.path.dirname(self.object_base_path)
+            child_prim_path = self.object_base_path
+            joint_prim_path = os.path.join(os.path.dirname(self.object_base_path), "FixedJoint")
+            joint_prim = stage.DefinePrim(joint_prim_path, "PhysicsFixedJoint")
+            joint_prim.GetRelationship("physics:body0").AddTarget(child_prim_path)
+            joint_prim.GetRelationship("physics:body1").AddTarget(parent_prim_path)
+
+    def get_observations(self):
+        translation, orientation = self.get_local_pose()
+        obs = {
+            "translation": translation,
+            "orientation": orientation,
+        }
+        return obs
+
+    def apply_texture(self, asset_root, cfg):
+        texture_name = cfg["texture_lib"]
+        texture_path_list = glob.glob(os.path.join(asset_root, texture_name, "*.jpg"))
+        texture_path_list.sort()
+        if cfg["apply_randomization"]:
+            texture_id = random.randint(0, len(texture_path_list) - 1)
+        else:
+            texture_id = cfg["texture_id"]
+        texture_path = texture_path_list[texture_id]
+        mat_prim_path = f"{self.base_prim_path}/Looks/Material"
+        mat = OmniPBR(
+            prim_path=mat_prim_path,
+            name="Material",
+            texture_path=texture_path,
+            texture_scale=cfg.get("texture_scale"),
+        )
+        self.apply_visual_material(mat)
+
+    def initialize(self):
+        super().initialize()
+        self._articulation_view.set_joint_velocities([0.0])
+        if "joint_position_range" in self.cfg:
+            self.articulation_initial_joint_position = np.random.uniform(
+                low=self.cfg["joint_position_range"][0], high=self.cfg["joint_position_range"][1]
+            )
+            self._articulation_view.set_joint_positions(
+                self.articulation_initial_joint_position, joint_indices=self.object_joint_index
+            )
+
+        if "strict_init" in self.cfg:
+            self._articulation_view.set_joint_position_targets(
+                self.cfg["strict_init"]["joint_positions"], joint_indices=self.cfg["strict_init"]["joint_indices"]
+            )
+        # self._articulation_view.set_joint_position_targets(self.target_joint_position)

workflows/simbox/core/objects/base_object.py

@@ -0,0 +1,8 @@
+OBJECT_DICT = {}
+
+
+def register_object(target_class):
+    key = target_class.__name__
+    # assert key not in OBJECT_DICT
+    OBJECT_DICT[key] = target_class
+    return target_class

workflows/simbox/core/objects/conveyor_object.py

@@ -0,0 +1,58 @@
+import os
+
+from core.objects.base_object import register_object
+from omni.isaac.core.prims import GeometryPrim
+from omni.isaac.core.utils.prims import create_prim
+from pxr import Gf, UsdPhysics
+
+
+@register_object
+class ConveyorObject(GeometryPrim):
+    def __init__(self, asset_root, root_prim_path, stage, cfg, *args, **kwargs):
+        """
+        Args:
+            asset_root: Asset root path
+            root_prim_path: Root prim path in USD stage
+            stage: USD stage
+            cfg: Config dict with required keys:
+                - name: Object name
+                - linear_velocity: [x, y, z] linear velocity in m/s
+                - linear_track_list: List of linear track names in USD
+                - angular_velocity: [x, y, z] angular velocity in rad/s
+                - angular_track_list: List of angular track names in USD
+        """
+        # ===== From cfg =====
+        self.asset_root = asset_root
+        self.stage = stage
+        prim_path = os.path.join(root_prim_path, cfg["name"])
+        usd_path = os.path.join(asset_root, cfg["path"])
+        self.linear_velocity = Gf.Vec3f(tuple(cfg["linear_velocity"]))
+        self.linear_track_list = cfg["linear_track_list"]
+        self.angular_velocity = Gf.Vec3f(tuple(cfg["angular_velocity"]))
+        self.angular_track_list = cfg["angular_track_list"]
+
+        # ===== Initialize =====
+        boxActorPath = prim_path
+        create_prim(usd_path=usd_path, prim_path=boxActorPath)
+        super().__init__(prim_path=boxActorPath, name=cfg["name"], *args, **kwargs)
+
+        # ===== Configure tracks (hard-coded logic) =====
+        linear_dir_list = [1, -1]
+        for i, linear_track in enumerate(self.linear_track_list):
+            belt_path = f"{boxActorPath}/World/{linear_track}/node_"
+            belt_prim = self.stage.GetPrimAtPath(belt_path)
+            UsdPhysics.CollisionAPI.Apply(belt_prim)
+            linarConveyor = UsdPhysics.RigidBodyAPI.Apply(belt_prim)
+            linarConveyor.CreateKinematicEnabledAttr().Set(True)
+            velocityAttribute = linarConveyor.GetVelocityAttr()
+            velocityAttribute.Set(linear_dir_list[i] * self.linear_velocity)
+
+        angular_dir = 1
+        for i, angular_track in enumerate(self.angular_track_list):
+            belt_path = f"{boxActorPath}/World/{angular_track}/validate_obj"
+            belt_prim = self.stage.GetPrimAtPath(belt_path)
+            UsdPhysics.CollisionAPI.Apply(belt_prim)
+            angularConveyor = UsdPhysics.RigidBodyAPI.Apply(belt_prim)
+            angularConveyor.CreateKinematicEnabledAttr().Set(True)
+            angularvelocityAttribute = angularConveyor.GetAngularVelocityAttr()
+            angularvelocityAttribute.Set(angular_dir * self.angular_velocity)

workflows/simbox/core/objects/geometry_object.py

@@ -0,0 +1,89 @@
+import glob
+import os
+import random
+
+from core.objects.base_object import register_object
+from omni.isaac.core.prims import GeometryPrim
+from omni.isaac.core.utils.prims import (
+    create_prim,
+    get_prim_at_path,
+    is_prim_path_valid,
+)
+
+try:
+    from omni.isaac.core.materials.omni_pbr import OmniPBR  # Isaac Sim 4.1.0 / 4.2.0
+except ImportError:
+    from isaacsim.core.api.materials import OmniPBR  # Isaac Sim 4.5.0
+
+
+@register_object
+class GeometryObject(GeometryPrim):
+    def __init__(self, asset_root, root_prim_path, cfg, *args, **kwargs):
+        """
+        Args:
+            asset_root: Asset root path
+            root_prim_path: Root prim path in USD stage
+            cfg: Config dict with required keys:
+                - name: Object name
+                - path: USD file path relative to asset_root
+                - prim_path_child (optional): Child prim path suffix
+        """
+        # ===== From cfg =====
+        self.asset_root = asset_root
+        prim_path = os.path.join(root_prim_path, cfg["name"])
+        usd_path = os.path.join(asset_root, cfg["path"])
+        if cfg.get("prim_path_child", None):
+            prim_path = os.path.join(prim_path, cfg["prim_path_child"])
+        self.cfg = cfg
+
+        # ===== Initialize =====
+        create_prim(prim_path=prim_path, usd_path=usd_path)
+        super().__init__(prim_path=prim_path, name=cfg["name"], *args, **kwargs)
+
+    def get_observations(self):
+        translation, orientation = self.get_local_pose()
+        scale = self.get_local_scale()
+        obs = {
+            "translation": translation,
+            "orientation": orientation,
+            "scale": scale,
+        }
+        return obs
+
+    def apply_texture(self, asset_root, cfg):
+        def _recursive_apply(prim, mat_prim_path):
+            for child in prim.GetChildren():
+                rel = child.GetRelationship("material:binding")
+                if rel:
+                    rel.SetTargets([mat_prim_path])
+                    continue
+                _recursive_apply(child, mat_prim_path)
+
+        texture_name = cfg["texture_lib"]
+        texture_path_list = glob.glob(os.path.join(asset_root, texture_name, "*"))
+        texture_path_list.sort()
+        if cfg["apply_randomization"]:
+            texture_id = random.randint(0, len(texture_path_list) - 1)
+        else:
+            texture_id = cfg["texture_id"]
+        texture_path = texture_path_list[texture_id]
+        mat_prim_path = f"{self.prim_path}/Looks/Material"
+        if not is_prim_path_valid(mat_prim_path):
+            self.mat = OmniPBR(
+                prim_path=mat_prim_path,
+                name="Material",
+                texture_path=texture_path,
+                texture_scale=cfg.get("texture_scale"),
+            )
+            target_prim_path = self.prim_path
+            if cfg.get("target_prim_path"):
+                target_prim_path = os.path.join(self.prim_path, cfg["target_prim_path"])
+                target_prim = get_prim_at_path(target_prim_path)
+            else:
+                target_prim = get_prim_at_path(target_prim_path)
+
+            _recursive_apply(target_prim, mat_prim_path)
+        else:
+            self.mat.set_texture(
+                texture_path,
+            )

workflows/simbox/core/objects/plane_object.py

@@ -0,0 +1,65 @@
+import glob
+import os
+import random
+
+from core.objects.base_object import register_object
+from omni.isaac.core.prims import XFormPrim
+from omni.isaac.core.utils.prims import is_prim_path_valid
+from omni.isaac.core.utils.stage import get_current_stage
+
+try:
+    from omni.isaac.core.materials.omni_pbr import OmniPBR  # Isaac Sim 4.1.0 / 4.2.0
+except ImportError:
+    from isaacsim.core.api.materials import OmniPBR  # Isaac Sim 4.5.0
+
+from pxr import UsdGeom
+
+
+@register_object
+class PlaneObject(XFormPrim):
+    def __init__(self, asset_root, root_prim_path, cfg, *args, **kwargs):
+        """
+        Args:
+            asset_root: Asset root path
+            root_prim_path: Root prim path in USD stage
+            cfg: Config dict with required keys:
+                - name: Object name
+                - size: [width, length] of the plane
+        """
+        # ===== From cfg =====
+        self.asset_root = asset_root
+        prim_path = os.path.join(root_prim_path, cfg["name"])
+        self.cfg = cfg
+
+        # ===== Initialize =====
+        stage = get_current_stage()
+        plane_geom = UsdGeom.Plane.Define(stage, prim_path)
+        plane_geom.CreateWidthAttr().Set(cfg["size"][0])
+        plane_geom.CreateLengthAttr().Set(cfg["size"][1])
+        super().__init__(prim_path=prim_path, name=cfg["name"], *args, **kwargs)
+
+    def get_observations(self):
+        raise NotImplementedError
+
+    def apply_texture(self, asset_root, cfg):
+        texture_name = cfg["texture_lib"]
+        texture_path_list = glob.glob(os.path.join(asset_root, texture_name, "*"))
+        texture_path_list.sort()
+        if cfg["apply_randomization"]:
+            texture_id = random.randint(0, len(texture_path_list) - 1)
+        else:
+            texture_id = cfg["texture_id"]
+        texture_path = texture_path_list[texture_id]
+        mat_prim_path = f"{self.prim_path}/Looks/Material"
+        if not is_prim_path_valid(mat_prim_path):
+            self.mat = OmniPBR(
+                prim_path=mat_prim_path,
+                name="Material",
+                texture_path=texture_path,
+                texture_scale=cfg.get("texture_scale"),
+            )
+            self.apply_visual_material(self.mat)
+        else:
+            self.mat.set_texture(
+                texture_path,
+            )

workflows/simbox/core/objects/rigid_object.py

@@ -0,0 +1,78 @@
+import glob
+import os
+import random
+
+from core.objects.base_object import register_object
+from omni.isaac.core.prims import RigidPrim
+from omni.isaac.core.utils.prims import create_prim, get_prim_at_path
+
+try:
+    from omni.isaac.core.materials.omni_pbr import OmniPBR  # Isaac Sim 4.1.0 / 4.2.0
+except ImportError:
+    from isaacsim.core.api.materials import OmniPBR  # Isaac Sim 4.5.0
+
+
+@register_object
+class RigidObject(RigidPrim):
+    def __init__(self, asset_root, root_prim_path, cfg, *args, **kwargs):
+        """
+        Args:
+            asset_root: Asset root path
+            root_prim_path: Root prim path in USD stage
+            cfg: Config dict with required keys:
+                - name: Object name
+                - path: USD file path relative to asset_root
+                - prim_path_child: Child prim path for rigid body
+                - init_translation (optional): Initial translation
+                - init_orientation (optional): Initial orientation
+                - init_parent (optional): Initial parent prim
+                - gap (optional): Gap parameter
+                - mass (optional): Object mass
+        """
+        # ===== From cfg =====
+        self.asset_root = asset_root
+        cfg_name = cfg["name"]
+        cfg_path = cfg["path"]
+        prim_path = f"{root_prim_path}/{cfg_name}"
+        usd_path = f"{asset_root}/{cfg_path}"
+        self.init_translation = cfg.get("init_translation", None)
+        self.init_orientation = cfg.get("init_orientation", None)
+        self.init_parent = cfg.get("init_parent", None)
+        self.gap = cfg.get("gap", None)
+        self.mass = cfg.get("mass", None)
+        kwargs["mass"] = cfg.get("mass", None)
+
+        # ===== Initialize =====
+        create_prim(prim_path=prim_path, usd_path=usd_path)
+        self.base_prim_path = prim_path
+        rigid_prim_path = os.path.join(self.base_prim_path, cfg["prim_path_child"])
+        self.mesh_prim_path = str(get_prim_at_path(rigid_prim_path).GetChildren()[0].GetPrimPath())
+        super().__init__(prim_path=rigid_prim_path, name=cfg["name"], *args, **kwargs)
+
+    def get_observations(self):
+        translation, orientation = self.get_local_pose()
+        scale = self.get_local_scale()
+        obs = {
+            "translation": translation,
+            "orientation": orientation,
+            "scale": scale,
+        }
+        return obs
+
+    def apply_texture(self, asset_root, cfg):
+        texture_name = cfg["texture_lib"]
+        texture_path_list = glob.glob(os.path.join(asset_root, texture_name, "*"))
+        texture_path_list.sort()
+        if cfg["apply_randomization"]:
+            texture_id = random.randint(0, len(texture_path_list) - 1)
+        else:
+            texture_id = cfg["texture_id"]
+        texture_path = texture_path_list[texture_id]
+        mat_prim_path = f"{self.base_prim_path}/Looks/Material"
+        mat = OmniPBR(
+            prim_path=mat_prim_path,
+            name="Material",
+            texture_path=texture_path,
+            texture_scale=cfg.get("texture_scale"),
+        )
+        self.apply_visual_material(mat)

workflows/simbox/core/objects/shape_object.py

@@ -0,0 +1,57 @@
+import glob
+import os
+import random
+
+from core.objects.base_object import register_object
+from omni.isaac.core.prims import GeometryPrim
+
+try:
+    from omni.isaac.core.materials.omni_pbr import OmniPBR  # Isaac Sim 4.1.0 / 4.2.0
+except ImportError:
+    from isaacsim.core.api.materials import OmniPBR  # Isaac Sim 4.5.0
+
+
+@register_object
+class ShapeObject(GeometryPrim):
+    def __init__(self, asset_root, root_prim_path, cfg, *args, **kwargs):
+        """
+        Args:
+            asset_root: Asset root path
+            root_prim_path: Root prim path in USD stage
+            cfg: Config dict with required keys:
+                - name: Object name
+        """
+        # ===== From cfg =====
+        self.asset_root = asset_root
+        prim_path = f"{root_prim_path}/{cfg.get('name')}"
+        self.cfg = cfg
+
+        # ===== Initialize =====
+        self.base_prim_path = prim_path
+        super().__init__(prim_path=self.base_prim_path, name=cfg.get("name"), *args, **kwargs)
+
+    def get_observations(self):
+        translation, orientation = self.get_local_pose()
+        obs = {
+            "translation": translation,
+            "orientation": orientation,
+        }
+        return obs
+
+    def apply_texture(self, asset_root, cfg):
+        texture_name = cfg.texture_lib
+        texture_path_list = glob.glob(os.path.join(asset_root, texture_name, "*.jpg"))
+        texture_path_list.sort()
+        if cfg.apply_randomization:
+            texture_id = random.randint(0, len(texture_path_list) - 1)
+        else:
+            texture_id = cfg.texture_id
+        texture_path = texture_path_list[texture_id]
+        mat_prim_path = f"{self.base_prim_path}/Looks/Material"
+        mat = OmniPBR(
+            prim_path=mat_prim_path,
+            name="Material",
+            texture_path=texture_path,
+            texture_scale=cfg.get("texture_scale"),
+        )
+        self.apply_visual_material(mat)

workflows/simbox/core/objects/xform_object.py

@@ -0,0 +1,65 @@
+import glob
+import os
+import random
+
+from core.objects.base_object import register_object
+from omni.isaac.core.prims import XFormPrim
+from omni.isaac.core.utils.prims import create_prim, is_prim_path_valid
+
+try:
+    from omni.isaac.core.materials.omni_pbr import OmniPBR  # Isaac Sim 4.1.0 / 4.2.0
+except ImportError:
+    from isaacsim.core.api.materials import OmniPBR  # Isaac Sim 4.5.0
+
+
+@register_object
+class XFormObject(XFormPrim):
+    def __init__(self, asset_root, root_prim_path, cfg, *args, **kwargs):
+        """
+        Args:
+            asset_root: Asset root path
+            root_prim_path: Root prim path in USD stage
+            cfg: Config dict with required keys:
+                - name: Object name
+                - path: USD file path relative to asset_root
+        """
+        # ===== From cfg =====
+        self.asset_root = asset_root
+        prim_path = os.path.join(root_prim_path, cfg["name"])
+        usd_path = os.path.join(asset_root, cfg["path"])
+        self.cfg = cfg
+
+        # ===== Initialize =====
+        create_prim(prim_path=prim_path, usd_path=usd_path)
+        super().__init__(prim_path=prim_path, name=cfg["name"], *args, **kwargs)
+
+    def get_observations(self):
+        translation, orientation = self.get_local_pose()
+        obs = {
+            "translation": translation,
+            "orientation": orientation,
+        }
+        return obs
+
+    def apply_texture(self, asset_root, cfg):
+        texture_name = cfg["texture_lib"]
+        texture_path_list = glob.glob(os.path.join(asset_root, texture_name, "*"))
+        texture_path_list.sort()
+        if cfg["apply_randomization"]:
+            texture_id = random.randint(0, len(texture_path_list) - 1)
+        else:
+            texture_id = cfg["texture_id"]
+        texture_path = texture_path_list[texture_id]
+        mat_prim_path = f"{self.prim_path}/Looks/Material"
+        if not is_prim_path_valid(mat_prim_path):
+            self.mat = OmniPBR(
+                prim_path=mat_prim_path,
+                name="Material",
+                texture_path=texture_path,
+                texture_scale=cfg.get("texture_scale"),
+            )
+            self.apply_visual_material(self.mat)
+        else:
+            self.mat.set_texture(
+                texture_path,
+            )