def make_env(idx, test): from pybullet_envs.bullet.kuka_diverse_object_gym_env import KukaDiverseObjectEnv # NOQA # Use different random seeds for train and test envs process_seed = int(process_seeds[idx]) env_seed = 2 ** 32 - 1 - process_seed if test else process_seed # Set a random seed for this subprocess misc.set_random_seed(env_seed) env = KukaDiverseObjectEnv( isDiscrete=True, renders=args.render and (args.demo or not test), height=84, width=84, maxSteps=max_episode_steps, isTest=test, ) # (84, 84, 3) -> (3, 84, 84) env = TransposeObservation(env, (2, 0, 1)) env = ObserveElapsedSteps(env, max_episode_steps) # KukaDiverseObjectEnv internally asserts int actions and does not # accept python-future's newint. env = CastAction(env, __builtins__.int) env.seed(int(env_seed)) if test and args.record: assert args.render,\ 'To use --record, --render needs be specified.' video_dir = os.path.join(args.outdir, 'video_{}'.format(idx)) os.mkdir(video_dir) env = RecordMovie(env, video_dir) return env
def make_env(idx, test): from pybullet_envs.bullet.kuka_diverse_object_gym_env import ( KukaDiverseObjectEnv, ) # NOQA # Use different random seeds for train and test envs process_seed = int(process_seeds[idx]) env_seed = 2 ** 32 - 1 - process_seed if test else process_seed # Set a random seed for this subprocess utils.set_random_seed(env_seed) env = KukaDiverseObjectEnv( isDiscrete=True, renders=args.render and (args.demo or not test), height=84, width=84, maxSteps=max_episode_steps, isTest=test, ) # Disable file caching to keep memory usage small env._p.setPhysicsEngineParameter(enableFileCaching=False) assert env.observation_space is None env.observation_space = gym.spaces.Box( low=0, high=255, shape=(84, 84, 3), dtype=np.uint8 ) # (84, 84, 3) -> (3, 84, 84) env = TransposeObservation(env, (2, 0, 1)) env = ObserveElapsedSteps(env, max_episode_steps) # KukaDiverseObjectEnv internally asserts int actions env = CastAction(env, int) env.seed(int(env_seed)) if test and args.record: assert args.render, "To use --record, --render needs be specified." video_dir = os.path.join(args.outdir, "video_{}".format(idx)) os.mkdir(video_dir) env = RecordMovie(env, video_dir) return env