Add SoftMimicGen for deformable objects

autodata_core/__init__.py

@@ -16,6 +16,7 @@
     GenerationAlgorithm,
     MimicGen,
     SkillGen,
+    SoftMimicGen,
     get_algorithm,
     iter_algorithms,
 )
@@ -33,6 +34,7 @@
     "GenerationResult",
     "MimicGen",
     "MultiWaypoint",
+    "SoftMimicGen",
     "SkillGen",
     "Waypoint",
     "WaypointSequence",

autodata_core/algorithms.py

@@ -7,21 +7,23 @@
 behavior diverges between Mimic, DexMimicGen, SkillGen, etc.
 Subclasses self-register at import time via :class:`_AlgorithmMeta`.
 
-The single behavioral hook is :meth:`GenerationAlgorithm.plan_subtask_trajectory`, called by
-``DataGenerator`` every time an EEF needs a new executable trajectory. The default implementation
-generates a subtask trajectory and merges an interpolation segment (the Mimic / DexMimicGen path).
-SkillGen overrides this to insert a motion-planned transit ahead of the merge.
+The main behavioral hook is :meth:`GenerationAlgorithm.plan_subtask_trajectory`, called by
+``DataGenerator`` every time an EEF needs a new executable trajectory. Object-state adaptation is
+routed through :meth:`GenerationAlgorithm.transform_source_eef_poses. Rigid algorithms use the
+default object-pose transform while SoftMimicGen overrides it with nodal TPS registration.
 """
 
 from __future__ import annotations
 
+import torch
 from collections.abc import Iterator
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     from autodata_core.data_generator import DataGenerator, _EEFGenerationState
     from autodata_core.waypoint import Waypoint
     from autodata_interfaces.datastream.datastream import Datastream
+    from autodata_interfaces.tasks.subtask_spec import Subtask
 
 REGISTERED_ALGORITHMS: dict[str, type[GenerationAlgorithm]] = {}
 
@@ -71,6 +73,45 @@ def validate_setup(self, datastream: Datastream) -> None:
         or ``datastream.get_env()`` (escape hatch) to enforce algorithm-specific invariants.
         """
 
+    def is_deformable_subtask(self, subtask: Subtask) -> bool:
+        """Return whether ``subtask`` uses deformable nodal state instead of a rigid pose."""
+
+        return False
+
+    def transform_source_eef_poses(
+        self,
+        *,
+        data_generator: DataGenerator,
+        eef_name: str,
+        subtask_ind: int,
+        subtask_object_name: str | None,
+        subtask_object_pose: torch.Tensor | None,
+        src_subtask_object_pose: torch.Tensor | None,
+        subtask_object_nodal_positions: torch.Tensor | None,
+        src_subtask_object_nodal_positions: torch.Tensor | None,
+        src_eef_poses: torch.Tensor,
+        use_delta_transform: torch.Tensor | None,
+        coord_transform_scheme: Any,
+        runtime_subtask_constraints_dict: dict,
+    ) -> torch.Tensor:
+        """Adapt source EEF poses to the current subtask object state.
+
+        The default implementation performs the existing rigid-object transform. Algorithms using
+        another object representation override this method.
+        """
+
+        return data_generator._apply_subtask_transform(
+            eef_name=eef_name,
+            subtask_ind=subtask_ind,
+            subtask_object_name=subtask_object_name,
+            subtask_object_pose=subtask_object_pose,
+            src_subtask_object_pose=src_subtask_object_pose,
+            src_eef_poses=src_eef_poses,
+            use_delta_transform=use_delta_transform,
+            coord_transform_scheme=coord_transform_scheme,
+            runtime_subtask_constraints_dict=runtime_subtask_constraints_dict,
+        )
+
     def plan_subtask_trajectory(
         self,
         *,
@@ -165,6 +206,96 @@ class DexMimicGen(GenerationAlgorithm):
     supports_coordination = True
 
 
+class SoftMimicGen(GenerationAlgorithm):
+    """MimicGen for deformable reference objects.
+
+    Deformable subtasks select source segments using nodal state and warp their EEF trajectories
+    with a thin-plate-spline transform. Rigid subtasks retain the normal MimicGen transform.
+    """
+
+    name = "softmimicgen"
+    expected_eef_count = (1, 2)
+    requires_motion_planner = False
+    uses_subtask_start_signals = False
+    supports_coordination = False
+
+    def is_deformable_subtask(self, subtask: Subtask) -> bool:
+        return bool(getattr(subtask.algo_params, "object_soft", False))
+
+    def validate_setup(self, datastream: Datastream) -> None:
+        live_nodal_positions = datastream.get_object_nodal_positions()
+        for eef_name in datastream.get_eef_names():
+            for subtask_index, subtask in enumerate(datastream.get_subtasks(eef_name)):
+                if not self.is_deformable_subtask(subtask):
+                    continue
+                assert (
+                    subtask.object_ref
+                ), f"SoftMimicGen deformable subtask {eef_name}[{subtask_index}] requires object_ref"
+                assert subtask.object_ref in live_nodal_positions, (
+                    f"Deformable object {subtask.object_ref!r} is not present in the live scene; "
+                    f"available: {sorted(live_nodal_positions)}"
+                )
+                live_node_count = live_nodal_positions[subtask.object_ref].shape[-2]
+                for demo_index, datagen_info in enumerate(datastream.source_pool.datagen_infos):
+                    assert (
+                        datagen_info.object_nodal_positions is not None
+                    ), f"Source demo {demo_index} lacks object_nodal_position annotations"
+                    assert (
+                        subtask.object_ref in datagen_info.object_nodal_positions
+                    ), f"Source demo {demo_index} lacks nodal positions for {subtask.object_ref!r}"
+                    source_node_count = datagen_info.object_nodal_positions[subtask.object_ref].shape[-2]
+                    assert source_node_count == live_node_count, (
+                        f"Source demo {demo_index} has {source_node_count} nodes for "
+                        f"{subtask.object_ref!r}, live object has {live_node_count}"
+                    )
+
+    def transform_source_eef_poses(
+        self,
+        *,
+        data_generator: DataGenerator,
+        eef_name: str,
+        subtask_ind: int,
+        subtask_object_name: str | None,
+        subtask_object_pose: torch.Tensor | None,
+        src_subtask_object_pose: torch.Tensor | None,
+        subtask_object_nodal_positions: torch.Tensor | None,
+        src_subtask_object_nodal_positions: torch.Tensor | None,
+        src_eef_poses: torch.Tensor,
+        use_delta_transform: torch.Tensor | None,
+        coord_transform_scheme: Any,
+        runtime_subtask_constraints_dict: dict,
+    ) -> torch.Tensor:
+        subtask = data_generator.datastream.get_subtask(eef_name, subtask_ind)
+        if not self.is_deformable_subtask(subtask):
+            return super().transform_source_eef_poses(
+                data_generator=data_generator,
+                eef_name=eef_name,
+                subtask_ind=subtask_ind,
+                subtask_object_name=subtask_object_name,
+                subtask_object_pose=subtask_object_pose,
+                src_subtask_object_pose=src_subtask_object_pose,
+                subtask_object_nodal_positions=subtask_object_nodal_positions,
+                src_subtask_object_nodal_positions=src_subtask_object_nodal_positions,
+                src_eef_poses=src_eef_poses,
+                use_delta_transform=use_delta_transform,
+                coord_transform_scheme=coord_transform_scheme,
+                runtime_subtask_constraints_dict=runtime_subtask_constraints_dict,
+            )
+
+        assert subtask_object_nodal_positions is not None, "current deformable nodal state is missing"
+        assert src_subtask_object_nodal_positions is not None, "source deformable nodal state is missing"
+        from autodata_core.deformable_transforms import transform_source_data_segment_using_nodal_registration
+
+        return transform_source_data_segment_using_nodal_registration(
+            src_eef_poses=src_eef_poses,
+            src_obj_nodal_pos=src_subtask_object_nodal_positions,
+            tgt_obj_nodal_pos=subtask_object_nodal_positions,
+            use_rotation_transform=subtask.algo_params.use_rotation_transform,
+            bend_coef=subtask.algo_params.bend_coef,
+            rot_coef=subtask.algo_params.rot_coef,
+        )
+
+
 class SkillGen(GenerationAlgorithm):
     """Single-arm SkillGen — motion-planned transit followed by skill replay per subtask.
 

autodata_core/data_generator.py

@@ -28,6 +28,7 @@
     transform_source_data_segment_using_object_pose,
 )
 from autodata_core.waypoint import MultiWaypoint, Waypoint, WaypointSequence, WaypointTrajectory
+from autodata_interfaces.env.reset_request import EnvResetRequest
 from autodata_interfaces.tasks.subtask_constraint_spec import SubTaskConstraintCoordinationScheme, SubTaskConstraintType
 
 if TYPE_CHECKING:
@@ -209,8 +210,10 @@ def select_source_demo(
         eef_name: str,
         eef_pose: torch.Tensor,
         object_pose: torch.Tensor | None,
+        object_nodal_positions: torch.Tensor | None,
         src_demo_current_subtask_boundaries: np.ndarray,
         subtask_object_name: str | None,
+        subtask_object_soft: bool,
         selection_strategy_name: str,
         selection_strategy_kwargs: dict | None = None,
     ) -> int:
@@ -228,7 +231,12 @@ def select_source_demo(
                     eef_pose=src_ep.eef_pose[eef_name][start_ind:end_ind],
                     object_poses=(
                         {subtask_object_name: src_ep.object_poses[subtask_object_name][start_ind:end_ind]}
-                        if subtask_object_name is not None
+                        if subtask_object_name is not None and not subtask_object_soft
+                        else None
+                    ),
+                    object_nodal_positions=(
+                        {subtask_object_name: src_ep.object_nodal_positions[subtask_object_name][start_ind:end_ind]}
+                        if subtask_object_name is not None and subtask_object_soft
                         else None
                     ),
                     subtask_term_signals=None,
@@ -243,6 +251,7 @@ def select_source_demo(
             eef_pose=eef_pose,
             object_pose=object_pose,
             src_subtask_datagen_infos=src_subtask_datagen_infos,
+            object_nodal_positions=object_nodal_positions,
             **kwargs,
         )
 
@@ -265,9 +274,15 @@ def generate_eef_subtask_trajectory(
         # Subtask.object_ref is empty string when no object is involved; normalize to None so the
         # rest of the pipeline can keep using the upstream `is not None` convention.
         subtask_object_name = subtasks[subtask_ind].object_ref or None
+        subtask_object_soft = self.algorithm.is_deformable_subtask(subtasks[subtask_ind])
         subtask_object_pose = (
             self.datastream.get_object_poses(env_ids=[env_id])[subtask_object_name][0]
-            if subtask_object_name is not None
+            if subtask_object_name is not None and not subtask_object_soft
+            else None
+        )
+        subtask_object_nodal_positions = (
+            self.datastream.get_object_nodal_positions(env_ids=[env_id])[subtask_object_name][0]
+            if subtask_object_name is not None and subtask_object_soft
             else None
         )
 
@@ -291,8 +306,10 @@ def generate_eef_subtask_trajectory(
                 eef_name=eef_name,
                 eef_pose=self.datastream.get_robot_eef_pose(env_ids=[env_id], eef_name=eef_name)[0],
                 object_pose=subtask_object_pose,
+                object_nodal_positions=subtask_object_nodal_positions,
                 src_demo_current_subtask_boundaries=all_randomized_subtask_boundaries[eef_name][:, subtask_ind],
                 subtask_object_name=subtask_object_name,
+                subtask_object_soft=subtask_object_soft,
                 selection_strategy_name=subtasks[subtask_ind].selection_strategy,
                 selection_strategy_kwargs=subtasks[subtask_ind].selection_strategy_kwargs,
             )
@@ -325,7 +342,14 @@ def generate_eef_subtask_trajectory(
             if channel_name == eef_name or channel_name not in eef_names
         }
         src_subtask_object_pose = (
-            src_ep.object_poses[subtask_object_name][selected_boundary[0]] if subtask_object_name is not None else None
+            src_ep.object_poses[subtask_object_name][selected_boundary[0]]
+            if subtask_object_name is not None and not subtask_object_soft
+            else None
+        )
+        src_subtask_object_nodal_positions = (
+            src_ep.object_nodal_positions[subtask_object_name][selected_boundary[0]]
+            if subtask_object_name is not None and subtask_object_soft
+            else None
         )
 
         if is_first_subtask or policy.transform_first_robot_pose:
@@ -343,12 +367,15 @@ def generate_eef_subtask_trajectory(
                 for channel_name, channel_tensor in src_subtask_passthrough_actions.items()
             }
 
-        transformed_eef_poses = self._apply_subtask_transform(
+        transformed_eef_poses = self.algorithm.transform_source_eef_poses(
+            data_generator=self,
             eef_name=eef_name,
             subtask_ind=subtask_ind,
             subtask_object_name=subtask_object_name,
             subtask_object_pose=subtask_object_pose,
             src_subtask_object_pose=src_subtask_object_pose,
+            subtask_object_nodal_positions=subtask_object_nodal_positions,
+            src_subtask_object_nodal_positions=src_subtask_object_nodal_positions,
             src_eef_poses=src_eef_poses,
             use_delta_transform=use_delta_transform,
             coord_transform_scheme=coord_transform_scheme,
@@ -623,8 +650,9 @@ async def _reset_environment_for_generation(
         # Recorder + reset queue stay on env. The initial scene state
         # snapshot is read through the Datastream interface.
         self.datastream.get_env().recorder_manager.reset(env_ids=env_id_tensor)
-        await env_reset_queue.put(env_id)
-        await env_reset_queue.join()
+        completion = asyncio.get_running_loop().create_future()
+        await env_reset_queue.put(EnvResetRequest(env_id=env_id, completion=completion))
+        await completion
         return env_id_tensor, self.datastream.get_scene_state(is_relative=True)
 
     def _build_runtime_subtask_constraints(self) -> dict:

autodata_core/datagen_info.py

@@ -22,6 +22,8 @@ class DatagenInfo:
     Attributes:
         eef_pose: ``{eef_name: [T, 4, 4]}`` recorded EEF poses [m, rad].
         object_poses: ``{object_name: [T, 4, 4]}`` recorded object poses [m, rad].
+        object_nodal_positions: ``{object_name: [T, N, 3]}`` recorded deformable-object
+            nodal positions [m].
         subtask_term_signals: ``{subtask_name: [T]}`` binary completion flag per step.
         subtask_start_signals: ``{subtask_name: [T]}`` binary start flag per step; required by SkillGen.
         target_eef_pose: ``{eef_name: [T, 4, 4]}`` controller target poses [m, rad].
@@ -34,13 +36,15 @@ def __init__(
         self,
         eef_pose: dict[str, torch.Tensor] | None = None,
         object_poses: dict[str, torch.Tensor] | None = None,
+        object_nodal_positions: dict[str, torch.Tensor] | None = None,
         subtask_term_signals: dict[str, Any] | None = None,
         subtask_start_signals: dict[str, Any] | None = None,
         target_eef_pose: dict[str, torch.Tensor] | None = None,
         passthrough_action: dict[str, torch.Tensor] | None = None,
     ) -> None:
         self.eef_pose = eef_pose
         self.object_poses = dict(object_poses) if object_poses is not None else None
+        self.object_nodal_positions = dict(object_nodal_positions) if object_nodal_positions is not None else None
         self.subtask_term_signals = dict(subtask_term_signals) if subtask_term_signals is not None else None
         self.subtask_start_signals = dict(subtask_start_signals) if subtask_start_signals is not None else None
         self.target_eef_pose = target_eef_pose
@@ -53,6 +57,8 @@ def to_dict(self) -> dict[str, Any]:
             out["eef_pose"] = self.eef_pose
         if self.object_poses is not None:
             out["object_poses"] = deepcopy(self.object_poses)
+        if self.object_nodal_positions is not None:
+            out["object_nodal_positions"] = deepcopy(self.object_nodal_positions)
         if self.subtask_start_signals is not None:
             out["subtask_start_signals"] = deepcopy(self.subtask_start_signals)
         if self.subtask_term_signals is not None: