data_generation.py

"""
Data generation for the case of Psm Envs and demonstrations.
Refer to
https://github.com/openai/baselines/blob/master/baselines/her/experiment/data_generation/fetch_data_generation.py
"""
import os
import argparse
import gym
import time
import numpy as np
import imageio
import cv2
from surrol.const import ROOT_DIR_PATH
from dsp.utils.general_utils import set_seed_everywhere
from dsp.surrol_wrappers import register_traj_imperfect_wrappers

parser = argparse.ArgumentParser(description='generate demonstrations for imitation')
parser.add_argument('--env', type=str, required=True,
                    help='the environment to generate demonstrations')
parser.add_argument('--max_episode_steps', type=int, default=50,
                    help='the max number of steps for the environment')
parser.add_argument('--num_episodes', type=int, default=100,
                    help='the number of episodes')
parser.add_argument('--video', action='store_true',
                    help='whether or not to record video')
parser.add_argument('--noise_strength', type=float, default=0.2,
                    help='strength of noise to be added')
parser.add_argument('--num_noisy_steps', type=int, default=5,
                    help='nums of noisy steps')
parser.add_argument('--use_success_only', action='store_true',
                    help='whether or not to use only successful episodes')
parser.add_argument('--seed', type=int, default=0,
                    help='random seed for the environment')
parser.add_argument('--noise_type', type=str, default='gaussian',
                    choices=['gaussian', 'uniform', 'poisson', 'salt_and_pepper'],
                    help='type of noise to be added to actions')
parser.add_argument('--msg', type=str, default='',
                    help='message to be added to the file name for saving data')
args = parser.parse_args()



actions = []
observations = []
infos = []

images = []  # record video
masks = []
EFFECTIVE_USE_SUCCESS_ONLY = False
EFFECTIVE_USE_FAILED_ONLY = False


def main():
    global EFFECTIVE_USE_SUCCESS_ONLY, EFFECTIVE_USE_FAILED_ONLY
    set_seed_everywhere(args.seed)
    register_traj_imperfect_wrappers()
    if args.noise_strength > 0: 
        print(f"using noise_type: {args.noise_type}")
    num_itr = args.num_episodes if not args.video else 5
    folder = 'demo' if not args.video else 'video'
    folder = os.path.join('datasets', folder)
    os.makedirs(folder) if not os.path.exists(folder) else None

    env = gym.make(args.env, render_mode='rgb_array')  # 'human', 'rgb_array'
    env._max_episode_steps  = args.max_episode_steps
    env.reset()  # trigger task setup before reading env-side success policy
    env_unwrapped = getattr(env, 'unwrapped', env)
    force_success_from_env = bool(getattr(env_unwrapped, 'force_success_for_traj_imperfect', False))
    force_failed_from_env = bool(getattr(env_unwrapped, 'force_failed_for_traj_imperfect', False))
    if force_success_from_env and force_failed_from_env:
        raise ValueError("Environment cannot force both success-only and failed-only collection.")
    effective_use_success_only = args.use_success_only or force_success_from_env
    effective_use_failed_only = force_failed_from_env and not effective_use_success_only
    EFFECTIVE_USE_SUCCESS_ONLY = effective_use_success_only
    EFFECTIVE_USE_FAILED_ONLY = effective_use_failed_only

    file_name = "data_"
    file_name += args.env
    file_name += "_" + str(num_itr)
    file_name += '_' + str(args.max_episode_steps)
    if args.num_noisy_steps > 0:
        file_name += f"_{args.noise_type}-noise"
        file_name += f"_0{int(args.noise_strength*10)}_{args.num_noisy_steps}"
    file_name += f"_seed{args.seed}"
    if args.msg:
        file_name += f"_{args.msg}"
    file_name += ".npz"


    if os.path.exists(os.path.join(folder, file_name)):
        print(f"\n############ DATASET EXIST ##########")
        print(f"File {file_name} already exists.\n")
        env.close()
        return

    cnt = 0
    init_time = time.time()

    if effective_use_success_only:
        print("Using only successful episodes.")
    elif effective_use_failed_only:
        print("Using only failed episodes.")
    # else:
    #     print("Using all episodes(including failed ones).")

    success_count = 0
    failed_count = 0

    while len(actions) < num_itr:
        obs = env.reset()
        print("ITERATION NUMBER ", len(actions))
        success = goToGoal(env, obs)
        cnt += 1

        success_count += 1 if success else 0
        failed_count += 0 if success else 1

    print(f"Success: {success_count}, Failed: {failed_count}")

    np.savez_compressed(os.path.join(folder, file_name),
                        acs=actions, obs=observations, info=infos)  # save the file

    if args.video:
        video_name = "video_"
        video_name += file_name[5:-4] + ".mp4"
        writer = imageio.get_writer(os.path.join(folder, video_name), fps=20)

        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = .8  # font size
        thickness = 2   # thickness of font
        color = (255, 255, 255)  # text color
        position = (10, 30)  # position of text (x, y)


        for img in images:
            img = img.copy()
            caption = f"perturbed training data: {args.env}"
            cv2.putText(img, caption, position, font, font_scale, color, thickness, lineType=cv2.LINE_AA)

            writer.append_data(img)
        writer.close()

        if len(masks) > 0:
            mask_name = "mask_"
            mask_name += args.env + ".npz"
            np.savez_compressed(os.path.join(folder, mask_name),
                                masks=masks)  # save the file

    used_time = time.time() - init_time
    print("Saved data at:", folder)
    print("Time used: {:.1f}m, {:.1f}s\n".format(used_time // 60, used_time % 60))
    print(f"Trials: {num_itr}/{cnt}")
    env.close()


def goToGoal(env, last_obs):
    episode_acs = []
    episode_obs = []
    episode_info = []

    time_step = 0  # count the total number of time steps
    episode_init_time = time.time()
    episode_obs.append(last_obs)

    obs, success = last_obs, False

    if args.num_noisy_steps > 0:
        num_noisy_step_ind = np.random.choice(args.max_episode_steps, args.num_noisy_steps).tolist()
    


    while time_step < env._max_episode_steps:
        ################### lookup table for act_mask ###################
        # for psm envs, we add noise to positions
        if args.env in {'NeedlePick-v0', 'NeedleReach-v0', 'PegTransfer-v0', 'GauzeRetrieve-v0'}:
            act_mask = np.array([1,1,1,0,0])
        elif args.env in {'BiPegTransfer-v0', 'NeedleRegrasp-v0'}:
            act_mask = np.array([1,1,1,0,0,1,1,1,0,0])

        # don't generate noise in ECM Tasks
        elif args.env in {'ECMReach-v0', 'StaticTrack-v0'}:
            act_mask = np.array([0,0,0]) 
        elif args.env in {'MisOrient-v0'}:
            act_mask = np.array([0])
            
        else:
            raise ValueError('Not implemented')
        #################################################################


        noise_type = args.noise_type

        if args.num_noisy_steps > 0 and time_step in num_noisy_step_ind:
            # action = (env.get_oracle_action(obs) + np.random.normal(args.noise_strength, 0.05, act_mask.shape[0])
            #           * act_mask * ((np.random.rand(act_mask.shape[0]) < 0.5)))
            # use_noisy_action = True
            if noise_type == 'gaussian':
                action = (env.get_oracle_action(obs) + np.random.normal(args.noise_strength, 0.05, act_mask.shape[0])
                          * act_mask * ((np.random.rand(act_mask.shape[0]) < 0.5)))
            elif noise_type == 'uniform':
                action = (env.get_oracle_action(obs) + np.random.uniform(-args.noise_strength, args.noise_strength, act_mask.shape[0])
                          * act_mask * ((np.random.rand(act_mask.shape[0]) < 0.5)))
            elif noise_type == 'poisson':
                action = (env.get_oracle_action(obs) + np.random.poisson(args.noise_strength, act_mask.shape[0])
                          * act_mask * ((np.random.rand(act_mask.shape[0]) < 0.5)))

            use_noisy_action = True
        else:
            action = env.get_oracle_action(obs)
            use_noisy_action = False
        # print(action.shape)
        if args.video:
            # img, mask = env.render('img_array')
            img = env.render('rgb_array')
            images.append(img)
            # masks.append(mask)

        obs, reward, done, info = env.step(action)
        info['noisy_action'] = use_noisy_action
        # print(f" -> obs: {obs}, reward: {reward}, done: {done}, info: {info}.")
        time_step += 1

        if isinstance(obs, dict) and info['is_success'] > 0 and not success:
            print("Timesteps to finish:", time_step)
            success = True

        episode_acs.append(action)
        episode_info.append(info)
        episode_obs.append(obs)
    print("Episode time used: {:.2f}s\n".format(time.time() - episode_init_time))

    should_store = (
        (success and EFFECTIVE_USE_SUCCESS_ONLY)
        or ((not success) and EFFECTIVE_USE_FAILED_ONLY)
        or (not EFFECTIVE_USE_SUCCESS_ONLY and not EFFECTIVE_USE_FAILED_ONLY)
    )

    if should_store:
        print("Success" if success else "Failed")
        actions.append(episode_acs)
        observations.append(episode_obs)
        infos.append(episode_info)
    
    return success



if __name__ == "__main__":
    print("\n###############################")
    print("##     generating dataset    ##")
    print("###############################\n")
    print(args)
    main()