Here as interactive gym to play with the robot, (I still need some help)

[raul-machine-learning]

First the good news:

This is an interactive gym where you can experiment with pre-trained policies to control the robot in real time.
Here is how to use it:

• Double-click on a body to select it.
• Ctrl + left drag applies a torque to the selected object, resulting in rotation.
• Ctrl + right drag applies a force to the selected object in the (x,z) plane, resulting in translation.
• Ctrl + Shift + right drag applies a force to the selected object in the (x,y) plane.

However, there are a few limitations:

• When you move the cubes, the robot doesn't seem to register the new positions and instead attempts to pick them up from their original locations.
• Only the environment lerobot/act_aloha_sim_insertion_human appears to work occasionally. The others either don't function at all or cause the program to crash due to missing attributes that haven't been implemented in the gym.

I'd really appreciate feedback/guidance from the repo maintainers on how to improve this snippet to support more environments and tasks.

file interactive_gym.py:

import gymnasium as gym
import mujoco
import mujoco.viewer
import torch
import importlib
from lerobot.policies.utils import get_device_from_parameters
from lerobot.configs import parser
from lerobot.configs.eval import EvalPipelineConfig
from lerobot.policies.factory import make_policy
from lerobot.envs.utils import preprocess_observation
from lerobot.utils.utils import get_safe_torch_device


# $ python interactive_gym.py --policy.path=lerobot/act_aloha_sim_insertion_human --env.type=aloha
# $ python interactive_gym.py --policy.path=lerobot/act_aloha_sim_transfer_cube_human --env.type=aloha

@parser.wrap()
def make_env_and_policy(cfg: EvalPipelineConfig):
    package_name = f"gym_{cfg.env.type}"

    try:
        importlib.import_module(package_name)
    except ModuleNotFoundError as e:
        print(f"{package_name} is not installed. Please install it with `pip install 'lerobot[{cfg.env.type}]'`")
        raise e

    gym_handle = f"{package_name}/{cfg.env.task}"

    env = gym.make(gym_handle, disable_env_checker=True, **cfg.env.gym_kwargs)

    policy = make_policy(cfg=cfg.policy, env_cfg=cfg.env)
    policy.eval()    
    policy.reset()

    return env, policy


def main(env, policy):
    device = get_device_from_parameters(policy)

    viewer = mujoco.viewer.launch_passive(env.unwrapped.model, env.unwrapped.data)

    observation, info = env.reset(seed=42)
    viewer.sync()

    for i in range(40000):
        observation = preprocess_observation(observation)
        observation = {
            key: observation[key].to(device, non_blocking=device.type == "cuda") for key in observation
        }

        # Infer "task" from attributes of environments.
        # TODO: works with SyncVectorEnv but not AsyncVectorEnv
        if hasattr(env, "task_description"):
            observation["task"] = env.unwrapped.task_description
        elif hasattr(env, "task"):
            observation["task"] = env.unwrapped.task
        else:  #  For envs without language instructions, e.g. aloha transfer cube and etc.
            observation["task"] = ""

        with torch.inference_mode():
            action = policy.select_action(observation)

        # Convert to CPU / numpy.
        action = action.to("cpu").numpy()
        assert action.ndim == 2, "Action dimensions should be (batch, action_dim)"

        # Apply the next action.
        #observation, reward, terminated, truncated, info = env.step(action)

        observation, reward, terminated, truncated, info = env.step(action[0])
        viewer.sync()
        
        if terminated or truncated:
            observation, info = env.reset()
            viewer.sync()
            
        if i % 100 == 0:
            print(i)

    viewer.close()
    env.close()

torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True

env, policy = make_env_and_policy()
main(env, policy)

[LokiMetaSmith]

hey, you might try using an LLM to help bounce code off of. people are great, but are super busy.

I threw your question at google gemini and got this back.

Why the Issues Occur
State Synchronization Lag: The main reason the robot doesn't "see" the moved cubes is due to how the simulation loop is structured. The observation variable, which contains the camera images and sensor data for the policy, is only updated when env.step() is called. When you move a cube with your mouse, you are directly changing the physics state in MuJoCo. However, the script only feeds the previous observation to the policy in the next loop iteration. The robot is essentially acting on information from one step in the past, before you moved the object.

Environment & Policy Incompatibility: The crashes you're seeing with other environments are almost certainly due to a mismatch between what the environment provides and what the policy expects. A policy is a bit like a specialized tool; one trained to perform a delicate insertion task with specific camera inputs (act_aloha_sim_insertion_human) won't know what to do with a different task (transfer_cube) that might have different goals or even different observation keys (e.g., a missing camera or a differently named sensor). The AttributeError happens when the policy code tries to access a key in the observation dictionary that the new environment doesn't provide.

Improved Script and Guidance
I've revised the script to be more robust, more interactive, and to provide you with the tools to diagnose these issues yourself. Below the code, I'll provide guidance on how to use it to tackle new environments.

Key Improvements in the New Script
Robust Main Loop: The for loop is replaced with while viewer.is_running():. This is the standard practice for mujoco.viewer and means the script will run until you close the visualization window, which feels more interactive.

Error Handling: The core logic inside the loop is now wrapped in a try...except AttributeError block. If you load an incompatible environment and the policy fails because it can't find an expected observation key, the program will no longer crash. Instead, it will print a helpful error message explaining the likely problem and exit the loop gracefully.

Debugging Information: When the script starts, it now prints a list of all keys in the environment's observation dictionary. This is your most powerful tool for debugging. If a policy for insertion expects wrist_camera_image but the transfer_cube environment provides main_camera_image, you will see the mismatch immediately.

Explanatory Comments: I've added comments directly in the code to explain why the state-sync issue happens and to clarify what each part of the loop does.

How to Move Forward and Support More Tasks
Diagnose with Observation Keys: When you try a new environment that fails, look at the "Initial observation keys" that are printed. Then, you'll need to find out what observation keys the policy expects. This information is often in the documentation for the model or can be inferred from its source code.

Create Environment Wrappers: The professional way to solve incompatibility is to write a "wrapper" for your gym environment. A wrapper is a class that sits between your script and the environment. It can modify the observation on the fly—for example, by renaming keys ('main_camera_image' -> 'wrist_camera_image') or even by removing keys the policy doesn't need. This is a more advanced topic but is the standard solution in reinforcement learning.

Understand Policy Specialization: Acknowledge that a policy is an expert at a single task. You can't expect an "insertion" policy to solve a "cube transfer" task. The goal of this script is to load a policy and evaluate it on the specific task it was trained for. The interactive controls are best used to introduce small disturbances or perturbations to see how the policy reacts and recovers, rather than to fundamentally change the task.

import gymnasium as gym
import mujoco
import mujoco.viewer
import torch
import importlib
import time
from lerobot.policies.utils import get_device_from_parameters
from lerobot.configs import parser
from lerobot.configs.eval import EvalPipelineConfig
from lerobot.policies.factory import make_policy
from lerobot.envs.utils import preprocess_observation
from lerobot.utils.utils import get_safe_torch_device

# To run, use commands like:
# $ python interactive_gym.py --policy.path=lerobot/act_aloha_sim_insertion_human --env.type=aloha
# $ python interactive_gym.py --policy.path=lerobot/act_aloha_sim_transfer_cube_human --env.type=aloha

@parser.wrap()
def make_env_and_policy(cfg: EvalPipelineConfig):
    """
    Initializes the gymnasium environment and the lerobot policy based on the provided configuration.
    """
    package_name = f"gym_{cfg.env.type}"

    try:
        importlib.import_module(package_name)
    except ModuleNotFoundError as e:
        print(f"Error: {package_name} is not installed.")
        print(f"Please install it with: pip install 'lerobot[{cfg.env.type}]'")
        raise e

    print(f"Making environment: {cfg.env.task}")
    gym_handle = f"{package_name}/{cfg.env.task}"
    env = gym.make(gym_handle, disable_env_checker=True, **cfg.env.gym_kwargs)

    print(f"Loading policy: {cfg.policy.path}")
    policy = make_policy(cfg=cfg.policy, env_cfg=cfg.env)
    policy.eval()
    policy.reset()

    return env, policy


def main(env, policy):
    """
    Runs the main interactive simulation loop.
    """
    device = get_device_from_parameters(policy)
    print(f"Running policy on device: {device}")

    try:
        # Launch the passive viewer. It allows interaction while the script runs.
        viewer = mujoco.viewer.launch_passive(env.unwrapped.model, env.unwrapped.data)
    except Exception as e:
        print(f"Error launching MuJoCo viewer: {e}")
        print("Please ensure you have a display environment (e.g., a desktop) available.")
        env.close()
        return

    # Reset the environment to get the initial observation.
    observation, info = env.reset(seed=42)
    
    # --- Helper for debugging ---
    # Print the keys of the observation dictionary to help diagnose mismatches.
    if isinstance(observation, dict):
        print(f"Initial observation keys: {list(observation.keys())}")
    # --------------------------

    # Use viewer.is_running() for a more interactive loop that ends when the window is closed.
    while viewer.is_running():
        start_time = time.time()
        
        # IMPORTANT: The state synchronization issue you observed is due to the nature of this loop.
        # User interactions in the viewer directly modify the physics state (env.unwrapped.data).
        # However, the `observation` variable used by the policy is only updated by `env.step()`.
        # This means the policy is always acting on the state from *before* your interaction in that step.
        # A true fix requires modifying the environment to refetch observations on demand, which is non-standard.
        # For now, this script treats user interaction as a "perturbation" to the policy's execution.

        try:
            # Preprocess the observation from the *previous* step.
            processed_obs = preprocess_observation(observation)
            processed_obs = {
                key: processed_obs[key].to(device, non_blocking=device.type == "cuda") for key in processed_obs
            }

            # Infer "task" from attributes of the environment.
            # This is a common point of failure if an env doesn't have these attributes.
            if hasattr(env.unwrapped, "task_description"):
                processed_obs["task"] = env.unwrapped.task_description
            elif hasattr(env.unwrapped, "task"):
                processed_obs["task"] = env.unwrapped.task
            else:
                processed_obs["task"] = ""

            # The policy selects an action based on the processed observation.
            with torch.inference_mode():
                action = policy.select_action(processed_obs)

            # Convert action to a CPU numpy array for the environment.
            action = action.to("cpu").numpy()
            assert action.ndim == 2, "Action dimensions should be (batch, action_dim)"

            # Apply the action to the environment and get the next observation.
            observation, reward, terminated, truncated, info = env.step(action[0])
            
            # Sync the viewer to reflect the new state.
            viewer.sync()

            # If the episode is over, reset the environment and the policy's internal state.
            if terminated or truncated:
                print("Episode finished. Resetting environment.")
                observation, info = env.reset()
                policy.reset()
                viewer.sync()

        except AttributeError as e:
            print("\n--- ATTRIBUTE ERROR ---")
            print(f"Caught an error: {e}")
            print("This usually means the policy is incompatible with the current environment.")
            print("Check the 'Initial observation keys' printed at the start.")
            print("The policy likely expects a key that this environment does not provide.")
            print("Stopping simulation loop.\n")
            break
        except Exception as e:
            print(f"\nAn unexpected error occurred: {e}")
            print("Stopping simulation loop.\n")
            break
            
        # Optional: Cap the loop frequency to avoid overwhelming the CPU.
        time.sleep(max(0, 0.01 - (time.time() - start_time)))


    viewer.close()
    env.close()
    print("Viewer closed. Exiting.")


if __name__ == "__main__":
    # These settings can improve performance on NVIDIA GPUs.
    torch.backends.cudnn.benchmark = True
    torch.backends.cuda.matmul.allow_tf32 = True

    # Parse arguments and initialize the environment and policy.
    env, policy = make_env_and_policy()
    # Run the main simulation.
    main(env, policy)

[tic80]

Hello, I'm encountering an error when trying to run your script viewer with the Aloha environment.

python interactive_gym.py --policy.path=lerobot/act_aloha_sim_insertion_human --env.type=aloha

When the script attempts to launch the MuJoCo viewer

viewer = mujoco.viewer.launch_passive(env.unwrapped.model, env.unwrapped.data)

It fails with the following error:

Error launching MuJoCo viewer: 'AlohaEnv' object has no attribute 'model'

I'm using

• MuJoCo version: 2.3.7
• Gymnasium 0.29.1
• gym-aloha 0.1.1

Do I need to use a different version ?

Thanks!

[raul-machine-learning]

Hello, I'm encountering an error when trying to run your script viewer with the Aloha environment.

Error launching MuJoCo viewer: 'AlohaEnv' object has no attribute 'model'

Now that you mention that I added the following properties to gym_aloha, I needed that for the launch_passive to work (I already filed this bug huggingface/gym-aloha#15)

here is the code for copy pasting it

    @property
    def model(self) -> mujoco.MjModel:
        return self._env.physics.model.ptr

    @property
    def data(self) -> mujoco.MjData:
        return self._env.physics.data.ptr