Add factory_assembly_task and the relevant task examples #285

yami007007 · 2025-12-12T10:02:44Z

Summary

Add a generic assembly_task (bring two objects into proximity) and corresponding task examples: peg insertion and gear meshing.

Detailed description

Why this change?
- The assembly task is a common pattern for bringing two generic objects into proximity.
- peg_insert and gear_mesh examples are added to demonstrate how to use this atomic task in practice.
What was changed?
- Added a new atomic task: assembly_task, plus a test script for this task.
- Added two new examples: peg_insert and gear_mesh.

Open issues

For the peg_insert task, the current peg and hole assets are not compatible with Pinocchio. Ideally, this configuration should be defined at the task environment level so that everyone uses the correct setup by default. However, an effective solution has not been found yet, so this is left to Vikram Ramasamy for follow-up (see the discussion on slack ).

Verification

Dataset Preparation

Please dowload the pre-recorded hdf5 files from link and place them in the datasets folder.

Test Commands

Peg insertion task example

python isaaclab_arena/examples/policy_runner.py \
        --policy_type zero_action \
        --disable_pinocchio \
        --environment isaaclab_arena_environments.tabletop_peginsert_environment:PegInsertEnvironment \
         peg_insert

python isaaclab_arena/scripts/imitation_learning/record_demos.py \
        --dataset_file datasets/franka_peginsert.hdf5 \
        --disable_pinocchio \
        --num_demos 1 \
        --environment isaaclab_arena_environments.tabletop_peginsert_environment:PegInsertEnvironment \
         peg_insert \
         --teleop_device keyboard

Since AVP teleop support is not available for Franka, I used an old HDF5 dataset with AbsEEF + binary gripper. Before running replay_demos.py, please update franka.py as follows, then you can use the original HDF5 dataset for demonstration replay.

Set use_relative_mode from True to False.
Set scale from 0.5 to 1.

Original implementation of arm_action in franka.py.

arm_action: ActionTermCfg = DifferentialInverseKinematicsActionCfg(
    asset_name="robot",
    joint_names=["panda_joint.*"],
    body_name="panda_hand",
    controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=True, ik_method="dls"),
    scale=0.5,
    body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.107]),
)

python isaaclab_arena/scripts/imitation_learning/replay_demos.py \
        --dataset_file datasets/franka_peginsert_10demos.hdf5 \
        --disable_pinocchio \
        --environment isaaclab_arena_environments.tabletop_peginsert_environment:PegInsertEnvironment \
        peg_insert

Gear meshing task example

python isaaclab_arena/examples/policy_runner.py \
        --policy_type zero_action \
        --environment isaaclab_arena_environments.tabletop_gearmesh_environment:GearMeshEnvironment \
         gear_mesh

python isaaclab_arena/scripts/imitation_learning/record_demos.py \
        --dataset_file datasets/franka_gearmesh.hdf5 \
        --num_demos 1 \
        --environment isaaclab_arena_environments.tabletop_gearmesh_environment:GearMeshEnvironment \
         gear_mesh \
         --teleop_device keyboard

For the same reason as above, please update the arm_action in franka.py accordingly.

python isaaclab_arena/scripts/imitation_learning/replay_demos.py \
        --dataset_file datasets/franka_gearmesh_6demos.hdf5 \
        --environment isaaclab_arena_environments.tabletop_gearmesh_environment:GearMeshEnvironment \
        gear_mesh

.gitignore

cvolkcvolk · 2025-12-18T10:49:16Z

isaaclab_arena/assets/object_library.py

+@register_asset
+class SmallGear(FactoryObject):
+    """
+    A small reference gear for gear mesh task.


Nit, this could be used by other tasks as well not just for "gear mesh task"? Suggest to update the comment

viiik-inside · 2025-12-18T16:05:30Z

isaaclab_arena/assets/background_library.py

+
+
+@register_asset
+class FactoryTableBackground(LibraryBackground):


Lets maybe call this the SeattleLabTable :)

viiik-inside · 2025-12-19T08:15:23Z

isaaclab_arena/assets/object_library.py

+    def __init__(self, prim_path: str | None = None, initial_pose: Pose | None = None):
+        super().__init__(prim_path=prim_path, initial_pose=initial_pose)
+
+        # Create factory-specific articulation configuration


I understand that this is being done so one can set other properties to the articulation. I think we should fix this directly where we generally create the articulations so we dont have extra functions.

viiik-inside · 2025-12-19T08:21:46Z

isaaclab_arena/tasks/factory_assembly_task.py

+from isaaclab_arena.terms.events import set_object_pose
+from isaaclab_arena.utils.cameras import get_viewer_cfg_look_at_object
+
+FRANKA_PANDA_FACTORY_HIGH_PD_CFG = FRANKA_PANDA_HIGH_PD_CFG.copy()


Do you think it would be nice to have this as a function ? Maybe we will also use this for other tasks?

Hi @viiik-inside, what do you mean by “having this as a function”? Do you mean wrapping all of this into a standalone function?
This is a reference implementation of franka.py in Isaac Lab. I've put these configs in franka.py of Arena embodiments. I'd like to know your suggestions on this.

from isaaclab_assets.robots.franka import FRANKA_PANDA_HIGH_PD_CFG FRANKA_PANDA_FACTORY_HIGH_PD_CFG = FRANKA_PANDA_HIGH_PD_CFG.copy() FRANKA_PANDA_FACTORY_HIGH_PD_CFG.spawn.activate_contact_sensors = True FRANKA_PANDA_FACTORY_HIGH_PD_CFG.spawn.rigid_props.disable_gravity = True FRANKA_PANDA_FACTORY_HIGH_PD_CFG.actuators["panda_shoulder"].stiffness = 150.0 FRANKA_PANDA_FACTORY_HIGH_PD_CFG.actuators["panda_shoulder"].damping = 30.0 FRANKA_PANDA_FACTORY_HIGH_PD_CFG.actuators["panda_forearm"].stiffness = 150.0 FRANKA_PANDA_FACTORY_HIGH_PD_CFG.actuators["panda_forearm"].damping = 30.0 FRANKA_PANDA_FACTORY_HIGH_PD_CFG.actuators["panda_hand"].stiffness = 150.0 FRANKA_PANDA_FACTORY_HIGH_PD_CFG.actuators["panda_hand"].damping = 30.0 FRANKA_PANDA_FACTORY_HIGH_PD_CFG.init_state.pos = (0.0, 0.0, 0.0) # for factory assembly task```

viiik-inside · 2025-12-19T08:30:24Z

isaaclab_arena/tasks/factory_assembly_task.py

+        self.held_asset = held_asset
+        self.assist_asset_list = assist_asset_list
+        self.background_scene = background_scene
+        self.scene_config = InteractiveSceneCfg(num_envs=1, env_spacing=3.0, replicate_physics=False)


Do we need this here? If yes, maybe lets make them parameters

viiik-inside · 2025-12-19T08:32:57Z

isaaclab_arena/tasks/factory_assembly_task.py

+            params={
+                "object_cfg": SceneEntityCfg(self.held_asset.name),
+                "target_object_cfg": SceneEntityCfg(self.fixed_asset.name),
+                "max_x_separation": 0.020,  # Tolerance for assembly alignment


Would be nice to have these as parameters. Maybe a dataclass ?

viiik-inside · 2025-12-19T08:34:26Z

isaaclab_arena_environments/mdp/__init__.py

+
+"""This sub-module contains the functions that are specific to the environment."""
+
+from isaaclab.envs.mdp import *  # noqa: F401, F403


Is there a reason why this folder is called mdp? Just curious

This is following the convention used in IsaacLab mdp. The events, observations, and terminations for the task example are placed in the mdp folder.

viiik-inside · 2025-12-19T08:35:28Z

isaaclab_arena_environments/tabletop_gearmesh_environment.py

+                func=randomize_object_serials_pose,
+                mode="reset",
+                params={
+                    "pose_range": {"x": (0.25, 0.6), "y": (-0.20, 0.20), "z": (0.0, 0.0), "yaw": (-1.0, 1.0)},


We can use our pose dataclass here

viiik-inside · 2025-12-19T08:36:16Z

isaaclab_arena_environments/tabletop_gearmesh_environment.py

+        from isaaclab_arena_environments.mdp.events import randomize_object_serials_pose
+
+        @configclass
+        class EventCfgGearMesh:


Can this be in the task class?

viiik-inside · 2025-12-19T08:37:05Z

isaaclab_arena_environments/tabletop_peginsert_environment.py

+        from isaaclab_arena_environments import mdp
+
+        @configclass
+        class EventCfgPegInsert:


If it has to be here, then we can add it outside this function, on top of the file maybe?

alexmillane

Looks good! Thank you for doing this!

We have requested a few changes.

One additional request is that we add a test for the FactoryAssemblyTask. Such a such could just start the task, then teleport objects into a success configuration and then check that things succeed.

alexmillane · 2025-12-19T09:35:53Z

isaaclab_arena/assets/background_library.py

+    object_min_z = -0.1
+
+    def __init__(self):
+        super().__init__(scale=(1.0, 1.0, 1.0))


Why are we passing the scale of 1.0 here? I would have thought that the default scale is 1.0? Can we remove this?

alexmillane · 2025-12-19T09:47:22Z

isaaclab_arena/assets/object_utils.py

+def create_factory_articulation_cfg(
+    prim_path: str,
+    usd_path: str,
+    scale: tuple[float, float, float],
+    mass: float,
+    rigid_props: sim_utils.RigidBodyPropertiesCfg = RIGID_BODY_PROPS_HIGH_PRECISION,
+    contact_offset: float = 0.005,
+    rest_offset: float = 0.0,
+) -> ArticulationCfg:


To echo what @viiik-inside said above, we need to expose this functionality through our Object class. I would suggest that we make a proposal next week, formulated as an MR against this one.

Modify the code based on the new MR.

alexmillane · 2025-12-19T10:01:54Z

isaaclab_arena_environments/mdp/events.py

+            if (
+                i == 0
+            ):  # set the "small_gear" and "large_gear" positions according the position of fixed_asset('gear_base')


This seems to make some assumptions about the order of the assets in the list. Is that correct? I would guess that it makes the assumption that the "fixed_asset" appears first in the list. Could we enforce this assumption using the function parameters? For example we could do something like:

def randomize_object_serials_pose( env: ManagerBasedEnv, env_ids: torch.Tensor, fixed_asset_cfg: SceneEntityCfg, asset_cfgs: list[SceneEntityCfg], min_separation: float = 0.0, pose_range: dict[str, tuple[float, float]] = {}, max_sample_tries: int = 5000, relative_asset_cfgs: SceneEntityCfg | None = None, ):

Such that the fixed_asset is passed separately. That way we remove assumptions about the order of assets in the list.

alexmillane · 2025-12-19T10:04:34Z

isaaclab_arena_environments/mdp/events.py

+            asset_cfg = asset_cfgs[i]
+            asset = env.scene[asset_cfg.name]
+
+            # Write pose to simulation
+            pose_tensor = torch.tensor([pose_list[i]], device=env.device)
+            positions = pose_tensor[:, 0:3] + env.scene.env_origins[cur_env, 0:3]
+            orientations = math_utils.quat_from_euler_xyz(pose_tensor[:, 3], pose_tensor[:, 4], pose_tensor[:, 5])
+
+            asset.write_root_pose_to_sim(
+                torch.cat([positions, orientations], dim=-1), env_ids=torch.tensor([cur_env], device=env.device)
+            )
+            asset.write_root_velocity_to_sim(
+                torch.zeros(1, 6, device=env.device), env_ids=torch.tensor([cur_env], device=env.device)
+            )


The code to get an asset and set it's pose and velocity as zero is repeated in this file. Here and below.

Suggestion to move it out to another function in this same file. Something like:

def set_pose_and_zero_velocity(env, asset_cfg, positions, orientations):

This will make this function more easily understood.

alexmillane · 2025-12-19T10:06:31Z

isaaclab_arena_environments/mdp/events.py

+def randomize_object_serials_pose(
+    env: ManagerBasedEnv,
+    env_ids: torch.Tensor,
+    asset_cfgs: list[SceneEntityCfg],
+    min_separation: float = 0.0,
+    pose_range: dict[str, tuple[float, float]] = {},
+    max_sample_tries: int = 5000,
+    relative_asset_cfgs: SceneEntityCfg | None = None,
+):


Suggestion to rename this function to something that indicates what it does. It's more specific than just randomizing poses. It randomizes poses, and then does something special with the relative assets.

Perhaps a docstring would help?

alexmillane · 2025-12-19T10:45:14Z

isaaclab_arena_environments/tabletop_gearmesh_environment.py

+        task = FactoryAssemblyTask(
+            fixed_asset=gear_base,
+            held_asset=medium_gear,
+            assist_asset_list=[small_gear, large_gear],
+            background_scene=background,
+        )
+        task.events_cfg = EventCfgGearMesh()


We replace the event_cfg in the underlying task. Do we ever use the original event_cfg defined in FactoryAssemblyTask?

alexmillane · 2025-12-19T10:49:59Z

isaaclab_arena/tasks/factory_assembly_task.py

+FRANKA_PANDA_FACTORY_HIGH_PD_CFG.init_state.pos = (0.0, 0.0, 0.0)  # for factory assembly task
+
+
+class FactoryAssemblyTask(TaskBase):


FactoryAssemblyTask -> AssemblyTask

There's nothing specific to factories here.

alexmillane · 2025-12-19T10:51:41Z

isaaclab_arena/tasks/factory_assembly_task.py

+    """
+    Factory assembly task where an object needs to be assembled with a base object, like peg insert, gear mesh, etc.
+    """


Suggestion to expand this docstring. From what I can tell, the goal in this task is to bring two objects into proximity. There are some additional (distractor?) objects in the scene (which are not involved in the task but are there to make things harder?).

If possible, could you expand the docstring to make that clear?

alexmillane · 2025-12-19T10:54:57Z

isaaclab_arena_environments/tabletop_gearmesh_environment.py

+            randomize_gear_positions = EventTerm(
+                func=randomize_object_serials_pose,
+                mode="reset",
+                params={
+                    "pose_range": {"x": (0.25, 0.6), "y": (-0.20, 0.20), "z": (0.0, 0.0), "yaw": (-1.0, 1.0)},
+                    "min_separation": 0.18,
+                    "asset_cfgs": [SceneEntityCfg("gear_base"), SceneEntityCfg("medium_gear")],
+                    "relative_asset_cfgs": [SceneEntityCfg("small_gear"), SceneEntityCfg("large_gear")],
+                },
+            )


Could we think of a way of incorperating this EventCfg in the FactoryAssemblyTask.

For example, the FactoryAssemblyTask could take an additional parameter randomization_mode where we select among a few different options for randomizing the assets.

What do you think? Is that possible?

Good suggestion, thanks. I’ve updated the code accordingly.

alexmillane · 2025-12-19T10:56:06Z

isaaclab_arena_environments/tabletop_peginsert_environment.py

+            randomize_peg_hole_positions = EventTerm(
+                func=randomize_object_pose,
+                mode="reset",
+                params={
+                    "pose_range": {
+                        "x": (0.2, 0.6),
+                        "y": (-0.20, 0.20),
+                        "z": (0.0, 0.0),
+                        "yaw": (-1.0, 1.0),
+                    },
+                    "min_separation": 0.1,
+                    "asset_cfgs": [SceneEntityCfg("peg"), SceneEntityCfg("hole")],
+                },
+            )


Same as my comment above: what do you think about incorperating this into FactoryAssemblyTask with an option for reset mode.

- Generalize FactoryAssemblyTask to AssemblyTask for generic base-object assembly - Centralize EventsCfg in the atomic task with randomization_mode for asset randomization variants - Simplify assembly asset initialization via spawn_cfg_addon and asset_cfg_addon - Add test script for the assembly task - Add docstrings for relevant functions - Use parameters instead of hard-coded values for flexibility

yami007007 · 2025-12-24T17:49:37Z

Hi @alexmillane @viiik-inside @cvolkcvolk, thanks for your review and suggestions.
I’ve refactored the assembly task based on your feedback:

Generalize FactoryAssemblyTask to AssemblyTask for generic base-object assembly
Centralize EventsCfg in the atomic task with randomization_mode for asset randomization variants
Simplify assembly asset initialization via spawn_cfg_addon and asset_cfg_addon
Add test script for the assembly task
Add docstrings for relevant functions
Use parameters instead of hard-coded values for flexibility

I won’t reply to every comment individually to keep the thread concise, but I’ve gone through all of your suggestions and applied the corresponding changes. Could you please take another look?

yami007007 requested a review from viiik-inside December 16, 2025 08:07

cvolkcvolk reviewed Dec 18, 2025

View reviewed changes

.gitignore Outdated Show resolved Hide resolved

cvolkcvolk reviewed Dec 18, 2025

View reviewed changes

viiik-inside reviewed Dec 18, 2025

View reviewed changes

viiik-inside reviewed Dec 19, 2025

View reviewed changes

alexmillane requested changes Dec 19, 2025

View reviewed changes

yami007007 added 2 commits December 24, 2025 23:50

Add factory_assembly_task and the relevant task examples

988c301

Rebase the assembly task to main.

87a2ac6

yami007007 force-pushed the feature/factory_assembly branch from 82ae523 to 87a2ac6 Compare December 24, 2025 16:09

yami007007 added 2 commits December 25, 2025 00:56



		@register_asset
		class FactoryTableBackground(LibraryBackground):


		"""This sub-module contains the functions that are specific to the environment."""

		from isaaclab.envs.mdp import * # noqa: F401, F403

		FRANKA_PANDA_FACTORY_HIGH_PD_CFG.init_state.pos = (0.0, 0.0, 0.0) # for factory assembly task


		class FactoryAssemblyTask(TaskBase):

Add factory_assembly_task and the relevant task examples #285

Are you sure you want to change the base?

Add factory_assembly_task and the relevant task examples #285

Uh oh!

Conversation

yami007007 commented Dec 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Detailed description

Open issues

Verification

Dataset Preparation

Test Commands

Peg insertion task example

Gear meshing task example

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yami007007 Dec 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alexmillane left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yami007007 commented Dec 24, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

yami007007 commented Dec 12, 2025 •

edited

Loading

yami007007 Dec 24, 2025 •

edited

Loading