def generate_data(self, verbose=False):
    '\n        Generates N square environments and trajectories ((observation, action) pairs)\n        for each environment\n\n        Params:\n            envs (int): number of environments to generate\n            steps (int): how many steps an agent initially takes in each environment\n            env_size (tuple): size of environment (should be something like (4,4), (9,9), etc...)\n            save (bool): whether or not to save the dataset\n        \n        Returns:\n            Dict of "environments, observations, actions", each corresponding to: \n                        environments: Array shape: (num_envs, env_size_x, env_size_y, num_categories), \n                        observations: Array shape: (num_envs, steps, num_categories),\n                        actions: Array shape: (num_envs, steps, 4)\n        '
    env_size = self.env_size
    if (self.num_categories == None):
        self.num_categories = (env_size[0] * env_size[1])
    self.environments = np.zeros((self.num_envs, env_size[0], env_size[1], self.num_categories))
    self.observations = np.zeros((self.num_envs, self.steps, self.num_categories))
    self.actions = np.zeros((self.num_envs, self.steps, 4))
    self.positions = np.zeros((self.num_envs, self.steps, 2))
    for i in range(self.num_envs):
        env = self.square_env()
        (obs, acts, pos) = self.trajectory_4way(env)
        self.environments[i] = env
        self.observations[i] = obs
        self.actions[i] = acts
        self.positions[i] = pos
    self.data = {'environments': self.environments, 'observations': self.observations, 'actions': self.actions, 'positions': self.positions}
    if self.save:
        name = os.path.join(self.data_root, 'four_way_graph.pickle')
        with open(name, 'wb') as handle:
            pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL)
Generates N square environments and trajectories ((observation, action) pairs)
for each environment

Params:
    envs (int): number of environments to generate
    steps (int): how many steps an agent initially takes in each environment
    env_size (tuple): size of environment (should be something like (4,4), (9,9), etc...)
    save (bool): whether or not to save the dataset

Returns:
    Dict of "environments, observations, actions", each corresponding to: 
                environments: Array shape: (num_envs, env_size_x, env_size_y, num_categories), 
                observations: Array shape: (num_envs, steps, num_categories),
                actions: Array shape: (num_envs, steps, 4)
def generate_data(self, verbose=False):
    '\n        Generates N square environments and trajectories ((observation, action) pairs)\n        for each environment\n\n        Params:\n            envs (int): number of environments to generate\n            steps (int): how many steps an agent initially takes in each environment\n            env_size (tuple): size of environment (should be something like (4,4), (9,9), etc...)\n            save (bool): whether or not to save the dataset\n        \n        Returns:\n            Dict of "environments, observations, actions", each corresponding to: \n                        environments: Array shape: (num_envs, env_size_x, env_size_y, num_categories), \n                        observations: Array shape: (num_envs, steps, num_categories),\n                        actions: Array shape: (num_envs, steps, 4)\n        '
    env_size = self.env_size
    if (self.num_categories == None):
        self.num_categories = (env_size[0] * env_size[1])
    self.environments = np.zeros((self.num_envs, env_size[0], env_size[1], self.num_categories))
    self.observations = np.zeros((self.num_envs, self.steps, self.num_categories))
    self.actions = np.zeros((self.num_envs, self.steps, 4))
    self.positions = np.zeros((self.num_envs, self.steps, 2))
    for i in range(self.num_envs):
        env = self.square_env()
        (obs, acts, pos) = self.trajectory_4way(env)
        self.environments[i] = env
        self.observations[i] = obs
        self.actions[i] = acts
        self.positions[i] = pos
    self.data = {'environments': self.environments, 'observations': self.observations, 'actions': self.actions, 'positions': self.positions}
    if self.save:
        name = os.path.join(self.data_root, 'four_way_graph.pickle')
        with open(name, 'wb') as handle:
            pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL)