Upload make_env.py with huggingface_hub

make_env.py

@@ -0,0 +1,119 @@
+import gymnasium
+from gymnasium import spaces
+import numpy as np
+
+# Define the environment class. There is a (12*12) grid, and the agent can move in four directions and collect randomly initialized items. The items are divided into 21 categories represented by int in range(21), and the number of each category is a multiple of 4.
+
+class GridWorldEnv(gymnasium.Env):
+    metadata = {'render.modes': ['human']}
+
+    def __init__(self):
+        super(GridWorldEnv, self).__init__()
+        self.grid_size = (12, 12)
+        self.num_categories = 21
+        self.action_space = spaces.Discrete(5)  # 0: down, 1: right, 2: up, 3: left, 4: collect
+        # self.action_space = spaces.Box(low=0, high=5, shape=(1,), dtype=np.int64)
+        # self.observation_space = spaces.Dict({
+        #     'grid': spaces.Box(low=-1 / 20, high=20 / 20, shape=self.grid_size, dtype=np.float64),
+        #     'loc': spaces.Box(low=0 / 12, high=11 / 12, shape=(2,), dtype=np.float64),
+        #     'bag': spaces.Box(low=0 / 4, high=4 / 4, shape=(self.num_categories, 1), dtype=np.float64)})
+        self.observation_space = spaces.Box(low=-1, high=20, shape=(144+2+self.num_categories, ), dtype=np.float64)
+        self.max_steps = 4 * 12 * 12
+        self.reset()
+
+    def reset(self, seed=None):
+        self.grid = self.initialize_grid()
+        self.loc = np.random.randint(0, self.grid_size[0], 2)
+        self.bag = np.zeros((self.num_categories, 1), dtype=int)
+        self.steps = 0
+        return self._get_obs(), self._get_info()
+
+    def initialize_grid(self):  # Randomly initialize the grid
+        grid = np.zeros(self.grid_size, dtype=int)
+        total_cells = self.grid_size[0] * self.grid_size[1]
+        remaining_cells = total_cells
+        category_counts = []
+        # Ensure each category has at least 4 items and is a multiple of 4
+        for _ in range(self.num_categories - 1):
+            count = np.random.randint(1, 3) * 4
+            if count > remaining_cells:
+                    count = remaining_cells
+            category_counts.append(count)
+            remaining_cells -= count
+        # Assign the remaining cells to the last category
+        category_counts.append(remaining_cells)
+        # Shuffle the positions and assign categories
+        positions = np.random.permutation(total_cells)
+        index = 0
+        for category, count in enumerate(category_counts):
+            for _ in range(count):
+                x, y = divmod(positions[index], self.grid_size[1])
+                grid[x, y] = category
+                index += 1
+
+        return grid
+
+    def step(self, action):
+        action = int(action)
+        self.steps += 1
+        reward = -0.1 - np.sum(self.bag) / (12 * 12)
+        if action == 0:  # down
+            self.loc[0] = min(self.loc[0] + 1, self.grid_size[0] - 1)
+        elif action == 1:  # right
+            self.loc[1] = min(self.loc[1] + 1, self.grid_size[1] - 1)
+        elif action == 2:  # up
+            self.loc[0] = max(self.loc[0] - 1, 0)
+        elif action == 3:  # left
+            self.loc[1] = max(self.loc[1] - 1, 0)
+        elif action == 4:  # collect
+            x, y = self.loc
+            if self.grid[x, y] == -1:
+                reward -= 2  # Penalty for collecting an empty cell
+            else:
+                self.bag[self.grid[x, y]] += 1
+                self.grid[x, y] = -1  # Remove the item from the grid
+
+        # 如果bag中有4个相同的物品，消除这4个物品，并获得奖励=+1
+        if np.any(self.bag == 4):
+            category = np.argmax(self.bag == 4)
+            self.bag[category] -= 4
+            reward += 1
+
+        done = bool(np.sum(self.grid) == -12 * 12)  # Done when all items are collected
+        if done:
+            reward += 100
+        truncated = bool(self.steps >= self.max_steps)  # Truncate the episode after a certain number of steps
+        if truncated:
+            reward = -3 * (np.sum(self.bag) + np.sum(self.grid != -1))  # Penalty for not collecting enough items
+        
+        # print(f'grid: {self.grid}')
+        return self._get_obs(), reward, done, truncated, self._get_info()
+    
+    def _get_obs(self):
+        # return {'grid': np.copy(self.grid) / 20, 'loc': np.copy(self.loc) / 12, 'bag': np.copy(self.bag) / 4}
+        return np.concatenate((self.grid.flatten(), self.loc, self.bag.flatten()))
+    
+    def _get_info(self):
+        return {'grid': np.copy(self.grid) / 20, 'loc': np.copy(self.loc) / 12, 'bag': np.copy(self.bag) / 4}
+    
+    def render(self, mode='rgb_array'):
+        if mode == 'human':
+            print("Grid:")
+            print(self.grid)
+            print("Loc:", self.loc)
+            print("Bag:", self.bag)
+
+    def close(self):
+        pass
+
+if __name__ == "__main__":
+    # Example usage
+    env = GridWorldEnv()
+    obs = env.reset()
+    env.render()
+    obs, reward, done, info = env.step(1)  # Move right
+    env.render()
+    obs, reward, done, info = env.step(0)  # Move down
+    env.render()
+    obs, reward, done, info = env.step(4)  # Collect item
+    env.render()