Stuck when using TeleopSystem and DataCollectionWrapper together

[Bailey-24]

Describe the bug
I can ran successfully in tests/test_data_colloection.py and omnigibson/examples/teleoperation/robot_teleoperate_demo.py, but I can't use DataCollectionWrapper with the teleopeation system.
I change the code in omnigibson/examples/teleoperation/robot_teleoperate_demo.py, in the following, but it stuck, I can't control the robot by keyboard or quest pro.

To Reproduce

1. Go to ' ran test_data_colloection.py'
2. add 'teleopeation system in test_data_colloection.py'
3. See error: stuck

And other try:

1. Go to ' ran robot_teleoperate_demo.py'
2. add 'DataCollectionWrapper in robot_teleoperate_demo.py'
3. See error: stuck

Expected behavior
Run successfully with TeleopSystem and DataCollectionWrapper together

Desktop (please complete the following information):

• OS: [e.g. Ubuntu 20.04]
• Isaac Sim Version [e.g. 4.1.0]
• OmniGibson Version [e.g. 1.1.0]

Additional context
My changed code:
robot_teleoperate_demo.py :

    cfg = dict(scene=scene_cfg, robots=[robot_cfg], objects=object_cfg)

    # Create the environment
    env = og.Environment(configs=cfg)
    env = DataCollectionWrapper(
        env=env,
        output_path=collect_hdf5_path,
        only_successes=False,
    )
    for i in range(2):
        env.reset()
    # update viewer camera pose
    og.sim.viewer_camera.set_position_orientation(position=[-0.22, 0.99, 1.09], orientation=[-0.14, 0.47, 0.84, -0.23])
    # Start teleoperation system
    robot = env.robots[0]

    # Initialize teleoperation system
    teleop_sys = TeleopSystem(config=teleop_config, robot=robot, show_control_marker=True)
    teleop_sys.start()

    # main simulation loop
    for _ in range(2000):
        action = teleop_sys.get_action(teleop_sys.get_obs())
        env.step(action)
    env.save_data()
    print("Data saved!")        

    # Shut down the environment cleanly at the end
    teleop_sys.stop()
    og.clear()

or in https://github.com/StanfordVL/OmniGibson/blob/main/tests/test_data_collection.py,

def test_data_collect_and_playback():
    from telemoma.configs.base_config import teleop_config
    from telemoma.utils.camera_utils import RealSenseCamera

    from omnigibson.utils.teleop_utils import TeleopSystem
    robot_name = choose_from_options(options=ROBOTS, name="robot")
    arm_teleop_method = choose_from_options(options=TELEOP_METHOD, name="robot arm teleop method")
    if robot_name != "FrankaPanda":
        base_teleop_method = choose_from_options(options=TELEOP_METHOD, name="robot base teleop method")
    else:
        base_teleop_method = "keyboard"  # Dummy value since FrankaPanda does not have a base
    # Generate teleop config
    teleop_config.arm_left_controller = arm_teleop_method
    teleop_config.arm_right_controller = arm_teleop_method
    teleop_config.base_controller = base_teleop_method
    teleop_config.interface_kwargs["keyboard"] = {"arm_speed_scaledown": 0.04}
    teleop_config.interface_kwargs["spacemouse"] = {"arm_speed_scaledown": 0.04}
    if arm_teleop_method == "vision" or base_teleop_method == "vision":
        teleop_config.interface_kwargs["vision"] = {"camera": RealSenseCamera()}

    cfg = {
        "env": {
            "external_sensors": [],
        },
        "scene": {
            "type": "InteractiveTraversableScene",
            "scene_model": "Rs_int",
            "load_object_categories": ["floors", "breakfast_table"],
        },
        "robots": [
            {
                "type": "Fetch",
                "obs_modalities": [],
            }
        ],
        # Task kwargs
        "task": {
            "type": "BehaviorTask",
            # BehaviorTask-specific
            "activity_name": "assembling_gift_baskets",
            "online_object_sampling": True,
        },
    }

    if og.sim is None:
        # Make sure GPU dynamics are enabled (GPU dynamics needed for cloth)
        gm.ENABLE_OBJECT_STATES = True
        gm.USE_GPU_DYNAMICS = True
        gm.ENABLE_TRANSITION_RULES = False
    else:
        # Make sure sim is stopped
        og.sim.stop()

    vr_data = "/home/pi/Documents/OmniGibson/vr_data"
    collect_hdf5_path = os.path.join(vr_data, "test_data_collection.hdf5")
    playback_hdf5_path = os.path.join(vr_data, "test_data_playback.hdf5")
    # Create the environment (wrapped as a DataCollection env)
    env = og.Environment(configs=cfg)
    robot = env.robots[0]
    env = DataCollectionWrapper(
        env=env,
        output_path=collect_hdf5_path,
        only_successes=False,
    )
    for _ in range(2):
        env.reset()

    # Initialize teleoperation system
    teleop_sys = TeleopSystem(config=teleop_config, robot=robot, show_control_marker=True)
    teleop_sys.start()
    
    # Record 2 episodes
    for i in range(2000):
        action = teleop_sys.get_action(teleop_sys.get_obs())
        env.step(action)

    # Save this data
    env.save_data()
    print("Data collection done")
    # Clear the sim
    og.clear(
        physics_dt=0.001,
        rendering_dt=0.001,
        sim_step_dt=0.001,
    )

How to solve

[CSCSX]

I'm also experiencing the same issue. Have you managed to resolve it? If so, could you please share the solution? Thank you so much!

[wensi-ai]

Hi, could you provide the full log including where the code is stuck? Thanks!

[Bailey-24]

my colletct data code is
collect.txt.

my log is
log.txt

I debug and found it stuck in the line # data_env.save_data().
But before save the data, i still can not control the robot.
It can't receive the control signal together with the data collection.

[wensi-ai]

Hi, apologies for the late response.

It seems like the problem is these 2 lines, which stop the GUI from refreshing. We will discuss among the team on a fix to this, but for now you can comment out these 2 lines and it should work. Another thing is you need to change line 152 of your script into action = teleop_sys.get_action(teleop_sys.get_obs()). I tested it on my end and the script works fine.

Let me know if you need further assistance.

[wensi-ai]

Update: fixed with this PR, feel free to try again.

[Bailey-24]

Thank you.
I can successfully collect the data.
But when I add DataPlaybackWrapper to collect playback data, it still have some problems.
My code is
collect.txt

1. bug is AssertionError: Object with name breakfast_table already exists in scene!.
   debug_log.txt
   I commented the code of asset to solve the problem.

2. KeyError: 'ag_obj_constraint_params'
   debug_log2.txt

So, I change the code to this,
test_data_collection2.txt

It can control the robot arm, save the collect and playback data, but I can't control the gripper, How to solve?

and a small issue is when i use the collect.txt to collect data. When the robotic arm approaches an object, it suddenly slows down and feels like there is a strong resistance, how to solve it.

[dexin-wang]

set show_control_marker=False may solve this.

[dexin-wang]

It seems that adding control_markers causes the step process to stall.

[Bailey-24]

I change the set show_control_marker=False, but still does not work.

The problem is same in the

bug is AssertionError: Object with name breakfast_table already exists in scene!.
debug_log.txt
I commented the code of asset to solve the problem.

KeyError: 'ag_obj_constraint_params'
debug_log2.txt

when I ran the collect.txt.

The gripper is still can not close in test_data_collection2.txt.

Could you please give a example full code, how to collect data and playback data. Thank you.