text: """
  Modified from robosuite example scripts.
  A script to collect a batch of human demonstrations that can be used
  to generate a learning curriculum (see `demo_learning_curriculum.py`).
  
  The demonstrations can be played back using the `playback_demonstrations_from_pkl.py`
  script.
  """
  import os
  import shutil
  import time
  import argparse
  import datetime
  import h5py
  from glob import glob
  import numpy as np
  import json
  
  import robosuite as suite
  from robosuite import load_controller_config
  from robosuite.wrappers import DataCollectionWrapper, VisualizationWrapper
  from robosuite.utils.input_utils import input2action
  
  import init_path
  from demonstration_tools import CustomDataCollectionWrapper
  from robosuite_task_zoo.environments.manipulation import *
  
  import cv2
  
  def collect_human_trajectory(env, device, arm, env_configuration, remove_directory=[]):
      """
      Use the device (keyboard or SpaceNav 3D mouse) to collect a demonstration.
      The rollout trajectory is saved to files in npz format.
      Modify the DataCollectionWrapper wrapper to add new fields or change data formats.
  
      Args:
          env (MujocoEnv): environment to control
          device (Device): to receive controls from the device
          arms (str): which arm to control (eg bimanual) 'right' or 'left'
          env_configuration (str): specified environment configuration
      """
  
      env.reset()
      # ID = 2 always corresponds to agentview
      env.render()
  
      is_first = True
  
      task_completion_hold_count = -1  # counter to collect 10 timesteps after reaching goal
      device.start_control()
  
      # Loop until we get a reset from the input or the task completes
      saving = True
      count = 0
      while True:
          count += 1
          # Set active robot
          active_robot = env.robots[0] if env_configuration == "bimanual" else env.robots[arm == "left"]
  
          # Get the newest action
          action, grasp = input2action(
              device=device,
              robot=active_robot,
              active_arm=arm,
              env_configuration=env_configuration
          )
  
          # If action is none, then this a reset so we should break
          if action is None:
              print("Break")
              saving = False
              break
  
          # Run environment step
          env.step(action)
          env.render()
  
          # Also break if we complete the task
          if task_completion_hold_count == 0:
              break
  
          # state machine to check for having a success for 10 consecutive timesteps
          if env._check_success():
              if task_completion_hold_count > 0:
                  task_completion_hold_count -= 1  # latched state, decrement count
              else:
                  task_completion_hold_count = 10  # reset count on first success timestep
          else:
              task_completion_hold_count = -1  # null the counter if there's no success
  
      print(count)
      # cleanup for end of data collection episodes
      if not saving:
          remove_directory.append(env.ep_directory.split('/')[-1])
      env.close()
      return saving
  
  def gather_demonstrations_as_hdf5(directory, out_dir, env_info, remove_directory=[]):
      """
      Gathers the demonstrations saved in @directory into a
      single hdf5 file.
  
      The strucure of the hdf5 file is as follows.
  
      data (group)
          date (attribute) - date of collection
          time (attribute) - time of collection
          repository_version (attribute) - repository version used during collection
          env (attribute) - environment name on which demos were collected
  
          demo1 (group) - every demonstration has a group
              model_file (attribute) - model xml string for demonstration
              states (dataset) - flattened mujoco states
              actions (dataset) - actions applied during demonstration
  
          demo2 (group)
          ...
  
      Args:
          directory (str): Path to the directory containing raw demonstrations.
          out_dir (str): Path to where to store the hdf5 file. 
          env_info (str): JSON-encoded string containing environment information,
              including controller and robot info
      """
  
      hdf5_path = os.path.join(out_dir, "demo.hdf5")
      f = h5py.File(hdf5_path, "w")
  
      # store some metadata in the attributes of one group
      grp = f.create_group("data")
  
      num_eps = 0
      env_name = None  # will get populated at some point
  
      for ep_directory in os.listdir(directory):
          print(ep_directory)
          if ep_directory in remove_directory:
              print("Skipping")
              continue
          state_paths = os.path.join(directory, ep_directory, "state_*.npz")
          states = []
          actions = []
  
          for state_file in sorted(glob(state_paths)):
              dic = np.load(state_file, allow_pickle=True)
              env_name = str(dic["env"])
  
              states.extend(dic["states"])
              for ai in dic["action_infos"]:
                  actions.append(ai["actions"])
                  
          if len(states) == 0:
              continue
  
          # Delete the first actions and the last state. This is because when the DataCollector wrapper
          # recorded the states and actions, the states were recorded AFTER playing that action.
          del states[-1]
          assert len(states) == len(actions)
  
          num_eps += 1
          ep_data_grp = grp.create_group("demo_{}".format(num_eps))
  
          # store model xml as an attribute
          xml_path = os.path.join(directory, ep_directory, "model.xml")
          with open(xml_path, "r") as f:
              xml_str = f.read()
          ep_data_grp.attrs["model_file"] = xml_str
  
          # write datasets for states and actions
          ep_data_grp.create_dataset("states", data=np.array(states))
          ep_data_grp.create_dataset("actions", data=np.array(actions))
  
      # write dataset attributes (metadata)
      now = datetime.datetime.now()
      grp.attrs["date"] = "{}-{}-{}".format(now.month, now.day, now.year)
      grp.attrs["time"] = "{}:{}:{}".format(now.hour, now.minute, now.second)
      grp.attrs["repository_version"] = suite.__version__
      grp.attrs["env"] = env_name
      grp.attrs["env_info"] = env_info
  
      f.close()
  
  
  if __name__ == "__main__":
      # Arguments
      parser = argparse.ArgumentParser()
      parser.add_argument(
          "--directory",
          type=str,
          default="demonstration_data",
      )
      parser.add_argument("--environment", type=str, default="PegInHoleEnv")
      parser.add_argument("--robots", nargs="+", type=str, default="Panda", help="Which robot(s) to use in the env")
      parser.add_argument("--config", type=str, default="single-arm-opposed",
                          help="Specified environment configuration if necessary")
      parser.add_argument("--arm", type=str, default="right", help="Which arm to control (eg bimanual) 'right' or 'left'")
      parser.add_argument("--camera", type=str, default="agentview", help="Which camera to use for collecting demos")
      parser.add_argument("--controller", type=str, default="OSC_POSE",
                          help="Choice of controller. Can be 'IK_POSE' or 'OSC_POSE'")
      parser.add_argument("--device", type=str, default="spacemouse")
      parser.add_argument("--pos-sensitivity", type=float, default=1.5, help="How much to scale position user inputs")
      parser.add_argument("--rot-sensitivity", type=float, default=1.5, help="How much to scale rotation user inputs")
      parser.add_argument("--num-demonstration", type=int, default=50, help="How much to scale rotation user inputs")
      parser.add_argument("--task", type=str, default=None, help="for multitask name")
      
      args = parser.parse_args()
  
  
      # Get controller config
      controller_config = load_controller_config(default_controller=args.controller)
  
      # Create argument configuration
      config = {
          "robots": args.robots,
          "controller_configs": controller_config,
      }
  
      # Check if we're using a multi-armed environment and use env_configuration argument if so
      if "TwoArm" in args.environment:
          config["env_configuration"] = args.config
  
      # Create environment
  
      if args.task is not None:
          config["task"] = args.task
      
      env = eval(args.environment)(
          **config,
          has_renderer=True,
          has_offscreen_renderer=False,
          render_camera=args.camera,
          ignore_done=True,
          use_camera_obs=False,
          reward_shaping=True,
          control_freq=20,
      )
  
      # Wrap this with visualization wrapper
      env = VisualizationWrapper(env)
  
      # Grab reference to controller config and convert it to json-encoded string
      env_info = json.dumps(config)
  
      # wrap the environment with data collection wrapper
      tmp_directory = "demonstration_data/{}_tmp/{}".format(args.environment, str(time.time()).replace(".", "_"))
      env = CustomDataCollectionWrapper(env, tmp_directory)
  
      # initialize device
      if args.device == "keyboard":
          from robosuite.devices import Keyboard
  
          device = Keyboard(pos_sensitivity=args.pos_sensitivity, rot_sensitivity=args.rot_sensitivity)
          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)
      elif args.device == "spacemouse":
          from robosuite.devices import SpaceMouse
  
          device = SpaceMouse(9583, 50734, pos_sensitivity=args.pos_sensitivity, rot_sensitivity=args.rot_sensitivity)
      else:
          raise Exception(
              "Invalid device choice: choose either 'keyboard' or 'spacemouse'."
          )
  
      # make a new timestamped directory
      t1, t2 = str(time.time()).split(".")
      new_dir = os.path.join(args.directory, "{}_pick_{}_place_{}_{}_{}".format(args.environment, args.pick_id, args.place_id, t1, t2))
      os.makedirs(new_dir)
  
      # collect demonstrations
  
      remove_directory = []
      i = 0
      while i < args.num_demonstration:
          print(i)
          saving = collect_human_trajectory(env, device, args.arm, args.config, remove_directory)
          if saving:
              print(remove_directory)
              gather_demonstrations_as_hdf5(tmp_directory, new_dir, env_info, remove_directory)
              i += 1
  