Data auto-generation part using a simple state machine

[Papaercold]

Summary

This PR introduces a state-machine–based data generation pipeline for the SO101 pick-orange task in LeIsaac.

The main implementation lives under scripts/environments/state_machine, where a scripted finite state machine is used to generate deterministic pick-and-place demonstrations. To integrate this pipeline with the existing teleoperation and replay infrastructure, a new teleoperation device type named so101_state_machine is added, with corresponding changes under source/leisaac.

The implementation intentionally follows the coding style and structure of
scripts/environments/teleoperation/teleop_se3_agent.py, and includes detailed comments for readability and maintainability.

---

Key Features

• State-machine–based data generation

  • Implemented under scripts/environments/state_machine
  • Designed to be deterministic, readable, and easy to extend

• New teleoperation device: so101_state_machine

  • Integrated into the existing teleoperation interface
  • Allows scripted policies to drive the environment using the same action pipeline as keyboard/gamepad devices

• Replay compatibility

  • Recorded datasets can be replayed using the existing
    scripts/environments/teleoperation/replay.py
  • No custom replay logic is required

---

Usage Examples

Generate data

python scripts/environments/state_machine/pick_orange.py \
  --dataset_file=./datasets/dataset_test.hdf5 \
  --task=LeIsaac-SO101-PickOrange-v0 \
  --num_envs=1 \
  --device=cuda \
  --enable_cameras \
  --record \
  --num_demos=1

Replay recorded data

python scripts/environments/teleoperation/replay.py \
  --dataset_file=./datasets/dataset_test.hdf5 \
  --task=LeIsaac-SO101-PickOrange-v0 \
  --num_envs=1 \
  --device=cuda \
  --enable_cameras \
  --select_episodes 1 \
  --replay_mode=action \
  --task_type=so101_state_machine

---

Observed Issue / Open Question

When running the state-machine–based data generation script, the robot gripper exhibits a brief downward drop at the beginning of each episode before stabilizing.

At first glance, this behavior appears to be gravity-related. However, gravity is explicitly disabled at spawn time for all relevant teleoperation devices, including the newly introduced state-machine device:

def use_teleop_device(self, teleop_device) -> None:
    self.task_type = teleop_device
    if teleop_device in ["keyboard", "gamepad", "so101_state_machine"]:
        self.scene.robot.spawn.rigid_props.disable_gravity = True

Given this configuration, it is unclear whether the observed initial drop is truly caused by gravity. Other potential factors may include controller or drive initialization behavior, insufficient drive stiffness during the first few simulation steps, or the absence of an explicit action warm-start when the episode begins.

I would appreciate feedback on the following questions:

• Is this type of initial transient expected when introducing a scripted state machine as a teleoperation device?
• Could this behavior be related to controller initialization, drive parameter settings, or action preprocessing rather than gravity itself?
• Are there recommended best practices in Isaac Lab / LeIsaac for avoiding such initial transients when using scripted or non-interactive teleoperation devices?

Any insights, suggestions, or references to similar patterns in existing tasks would be greatly appreciated.

---

Notes

• This PR focuses on clarity, determinism, and minimal intrusion into existing teleoperation and replay logic.
• The issue described above does not affect replay correctness once the episode is running, but it does impact the visual behavior at the very beginning of an episode.

[EverNorif]

@Papaercold Thank you for your PR. Due to a busy schedule recently, I may not be able to review the code right away, I'll check it later.

Regarding the brief drop of the gripper that you mentioned, I haven’t observed this behavior so far. I’ll check for it again using the latest code.

[Papaercold]

@EverNorif
I have refactored the state machine implementation.

The state machine code has been moved to:

source/leisaac/leisaac/state_machine

The scripts for automatic data generation using the state machine are now located in:

scripts/environments/state_machine

In addition, I have fixed the issue where the robot arm would drop at the beginning due to gravity (i.e., it was not properly applying control effort during initialization).

Going forward, all task-specific state machines will be rewritten based on the StateMachineBase abstract class to ensure a consistent and extensible structure.

To make the review process easier, I can also provide a brief documentation file explaining the automatic data generation workflow and the structure of the state machine.

Please let me know if you would like me to add that documentation as part of this PR.

[Papaercold]

In addition, I would appreciate it if you could let me know your next plans.

Since the current structure has been refactored to be more modular and extensible, beyond the existing automatic pick-orange task, I can also add more tasks if needed, or integrate an RL module based on new requirements.

Thank you very much for your time and review.

Zihan Gao

[Papaercold]

I have now added STATE_MACHINE_README.md to the repository, including both English and Simplified Chinese versions.
Please feel free to refer to it during the review process.
Thank you very much for your help.

[Papaercold]

The fold_cloth state machine is still under development.

At the moment, it is very difficult to complete the cloth-folding task using the current state machine.

I believe there are still some issues with the collision configuration. In some cases, the gripper penetrates the cloth mesh, which prevents it from being properly grasped and lifted.

[EverNorif]

@Papaercold Thank you! I will review it next week.

[Papaercold]

I’ve started working on the reinforcement learning module based on the rsl_rl library, but it is still under development.

For now, please ignore the files under the rl folder as well as the new rl_config settings during your review.

The directory structure and corresponding file descriptions are also documented in the Markdown file.

If you prefer, I can submit a separate PR for the RL part after you finish reviewing the state machine implementation.

[EverNorif]

@Papaercold Yes, I think it would be better to keep the PR focused on a single feature. You can save your progress locally or in another branch first. This PR should only track the state machine–related functionality.

[Papaercold]

@Papaercold Yes, I think it would be better to keep the PR focused on a single feature. You can save your progress locally or in another branch first. This PR should only track the state machine–related functionality.

Sure, I’ll make the changes now. I’ll keep this PR focused on the state machine functionality.

[Papaercold]

@EverNorif The modifications have been completed. In this commit, the RL module has been fully removed. I have also removed fold_cloth.py, as its current performance is not yet satisfactory.

The remaining code consists only of the components that I have tested and verified to be functioning correctly. You may refer to the Markdown documentation for guidance during your review.

[EverNorif]

I tested this PR locally and it runs successfully.

There are still some issues with the structure and implementation that need to be adjusted.

In terms of structure, I think this organization would be better：

scripts/datagen/state_machine/
-- generate.py(or any other name you think is appropriate.)      # Unified runner script: run the state machine based on the provided task, rather than being limited to specific one.
-- replay.py           # Replay script for state-machine demonstrations

source/leisaac/leisaac/datagen/state_machine/
-- base.py             # StateMachineBase abstract class
-- pick_orange.py      # PickOrangeStateMachine
-- ... # other state machine implement

[EverNorif]

It is recommended to use a fixed version of IsaacLab rather than the latest one.

[EverNorif]

At the moment, it seems that each file corresponds to a single task. I think it would be better to refactor this into a more general entry point, generate.py, which takes different Tasks as input and generates the corresponding action.

In generate.py, we would only call the relevant interfaces from StateMachineBase, while the specific implementations would be placed in separate task-specific state machine modules.

[EverNorif]

This line should also be removed.

[EverNorif]

translate it to english

[EverNorif]

There’s no need to provide the launch script here, it can simply be documented instead.

[EverNorif]

There’s no need to provide the launch script here, it can simply be documented instead.

[EverNorif]

Later, this document can be moved to the docs folder and included as part of the website content.

[EverNorif]

Why is demo_0 always empty here? The teleop script doesn’t seem to have this issue. As far as I remember, if no data is recorded, reset() does not log any data.

[EverNorif]

You can provide more detailed error messages in the code.

[Papaercold]

@EverNorif Thanks for the review. I’ll make the requested changes shortly and follow up here if anything comes up.