Running demonstration with keyboard control

[pull-ups]

Hi, Thank you for sharing awesome codbase.
I have a question for crafting human demonstration data.

in scripts/collect_demonstration.py, it seems that we have 2 choice for collecting human demonstrations: keyboard and spacemouse.
However, I don't have spacemouse, so when I run

python scripts/collect_demonstration.py --device keyboard --robots Panda --bddl-file $BDDL_FILE_DIR
(BDDL_FILE_DIR is one of bddl file in directory.)

below error occurs:

Traceback (most recent call last):
  File "scripts/collect_demonstration.py", line 317, in <module>
    env.viewer.add_keypress_callback("any", device.on_press)
TypeError: add_keypress_callback() takes 2 positional arguments but 3 were given

When I change as below,

        # env.viewer.add_keypress_callback("any", device.on_press)
        # env.viewer.add_keyup_callback("any", device.on_release)
        # env.viewer.add_keyrepeat_callback("any", device.on_press)
        env.viewer.add_keypress_callback(device.on_press)
        env.viewer.add_keyup_callback(device.on_release)
        env.viewer.add_keyrepeat_callback(device.on_press)

And then below error occurs:

Traceback (most recent call last):
  File "scripts/collect_demonstration.py", line 321, in <module>
    env.viewer.add_keyup_callback(device.on_release)
AttributeError: 'OpenCVRenderer' object has no attribute 'add_keyup_callback'

There was no add_keyup_callback on Robosuit's OpenCVRender in the first place.

How can I solve it?
Thanks.

[n3puiol]

Hi, I encountered a similar issue. Seems like the code was written for an older robosuite version. Anyway I made an alternative demonstration collector.

import os
import argparse

import numpy as np
import torch
from matplotlib import pyplot as plt
from tqdm import tqdm
import imageio

from libero.libero import benchmark, get_libero_path
from libero.libero.envs import OffScreenRenderEnv

from pathlib import Path


class KeyboardController:
    """Handles keyboard input for robot control."""

    def __init__(self):
        self.move_speed = 0.8
        self.rot_speed = 0.2
        self.current_action = [0.0] * 7  # [x, y, z, roll, pitch, yaw, gripper]
        self.gripper_open = True

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        pass

    def get_key(self):
        """Get keyboard input (blocking using input())."""
        try:
            key = input("Enter command (or 'quit' to exit): ").strip().lower()
            if key == 'quit':
                return '\x1b'  # ESC equivalent
            elif key:
                return key[0]  # Take first character
            return None
        except (EOFError, KeyboardInterrupt):
            return '\x1b'

    def get_action_from_key(self, key):
        """
        Convert keyboard input to robot action.

        Movement controls:
        - w/s: forward/backward (x-axis)
        - a/d: left/right (y-axis)
        - q/e: up/down (z-axis)
        - i/k: pitch up/down
        - j/l: yaw left/right
        - u/o: roll left/right
        - g: toggle gripper
        - r: reset all movements
        """
        action = [0.0] * 7

        if key is None:
            return action

        # Position controls
        if key == 'w':
            action[0] = self.move_speed  # forward
        elif key == 's':
            action[0] = -self.move_speed  # backward
        elif key == 'a':
            action[1] = self.move_speed  # left
        elif key == 'd':
            action[1] = -self.move_speed  # right
        elif key == 'q':
            action[2] = self.move_speed  # up
        elif key == 'e':
            action[2] = -self.move_speed  # down

        # Orientation controls
        elif key == 'u':
            action[3] = self.rot_speed  # roll left
        elif key == 'o':
            action[3] = -self.rot_speed  # roll right
        elif key == 'i':
            action[4] = self.rot_speed  # pitch up
        elif key == 'k':
            action[4] = -self.rot_speed  # pitch down
        elif key == 'j':
            action[5] = self.rot_speed  # yaw left
        elif key == 'l':
            action[5] = -self.rot_speed  # yaw right

        # Gripper control
        elif key == 'g':
            self.gripper_open = not self.gripper_open
            print(f"Gripper {'opened' if self.gripper_open else 'closed'}")

        # Reset
        elif key == 'r':
            action = [0.0] * 6
            print("Position reset")

        # Set gripper state
        action[6] = -1.0 if self.gripper_open else 1.0

        return action

    def print_controls(self):
        """Print control instructions."""
        print("\n" + "="*50)
        print("KEYBOARD CONTROLS:")
        print("="*50)
        print("Type commands followed by ENTER.")
        print("Type 'quit' to exit.\n")
        print("Movement:")
        print("  W/S - Forward/Backward")
        print("  A/D - Left/Right")
        print("  Q/E - Up/Down")
        print("\nOrientation:")
        print("  U/O - Roll Left/Right")
        print("  I/K - Pitch Up/Down")
        print("  J/L - Yaw Left/Right")
        print("\nOther:")
        print("  G - Toggle Gripper")
        print("  R - Reset Position")
        print("  ESC/quit - Quit")
        print("="*50)


def main():
    parser = argparse.ArgumentParser(description="Run LIBERO with keyboard control")
    parser.add_argument("--task_suite_name", type=str, default="libero_goal",
                        help="Task suite name (e.g., libero_goal, libero_spatial, libero_object)")
    parser.add_argument("--task_id", type=int, default=2,
                        help="Task ID to evaluate")
    parser.add_argument("--seed", type=int, default=0,
                        help="Random seed")
    parser.add_argument("--steps", type=int, default=300,
                        help="Number of evaluation steps")
    parser.add_argument("--fps", type=int, default=20,
                        help="FPS for output video")
    parser.add_argument("--render", action="store_true", default=True, help="Render the environment")

    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "mps" if torch.mps.is_available() else "cpu"

    output_directory = Path(f"outputs/eval/{args.task_suite_name}")
    output_directory.mkdir(parents=True, exist_ok=True)

    benchmark_dict = benchmark.get_benchmark_dict()
    task_suite = benchmark_dict[args.task_suite_name]()

    # retrieve a specific task
    task = task_suite.get_task(args.task_id)
    task_name = task.name
    task_description = task.language
    task_bddl_file = os.path.join(get_libero_path("bddl_files"), task.problem_folder, task.bddl_file)
    print(f"[info] retrieving task {args.task_id} from suite {args.task_suite_name}, the " + \
          f"language instruction is {task_description}, and the bddl file is {task_bddl_file}")

    # step over the environment
    env_args = {
        "bddl_file_name": task_bddl_file,
        "camera_heights": 256,
        "camera_widths": 256,
    }
    env = OffScreenRenderEnv(**env_args)
    env.seed(args.seed)

    obs = env.reset()
    init_states = task_suite.get_task_init_states(
        args.task_id)  # for benchmarking purpose, we fix the a set of initial states
    init_state_id = 0
    env.set_init_state(init_states[init_state_id])

    frames = []

    # Initialize keyboard controller
    with KeyboardController() as controller:
        controller.print_controls()

        for step in tqdm(range(args.steps)):
            if step < 10:
                obs, reward, done, info = env.step([0.0] * 6 + [-1.0])
                continue

            # Get keyboard input and generate action
            key = controller.get_key()
            if key == '\x1b':  # ESC key
                print("\nExiting...")
                break

            action = controller.get_action_from_key(key)

            image = torch.tensor(obs["agentview_image"]).to(device, dtype=torch.uint8).permute(2, 0, 1) / 255.0
            image = torch.flip(image, dims=[1, 2])

            wrist_image = torch.tensor(obs["robot0_eye_in_hand_image"]).to(device, dtype=torch.uint8).permute(2, 0,
                                                                                                              1) / 255.0
            wrist_image = torch.flip(wrist_image, dims=[1, 2])

            if args.render:
                plt.figure(figsize=(10, 5))
                plt.subplot(1, 2, 1)
                plt.imshow(image.permute(1, 2, 0).cpu().numpy())
                plt.title("Agent View")
                plt.axis('off')

                plt.subplot(1, 2, 2)
                plt.imshow(wrist_image.permute(1, 2, 0).cpu().numpy())
                plt.title("Wrist Camera")
                plt.axis('off')

                plt.show()

            obs, reward, done, info = env.step(action)

            frame = (wrist_image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
            frames.append(frame)

    env.close()

    video_path = output_directory / f"{task_name}_keyboard_control.mp4"

    imageio.mimsave(str(video_path), frames, fps=args.fps)
    print(f"[info] Video of the evaluation is available in '{video_path}'.")


if __name__ == "__main__":
    main()

Hopefully this helps!

[pull-ups]

Dear Louis Vanstappen, Thanks a lot for sharing your awesome code!
I have ran your code and got a question for this.
When I run your code, the environment was not an real-time interactive environment. Rather, I could see the image after I press a keyboard, and again. Is this right you intend, or is there any other way to use this interactively?

Thanks.
Seungwon

[Boltzmachine]

I just delete the three lines. Looks fine so far