text: from gym_minigrid.minigrid import *
  from gym_minigrid.multigrid import *
  from gym_minigrid.register import register
  
  class Agent(WorldObj):
      def __init__(self, index=0):
          super(Agent, self).__init__('agent', IDX_TO_COLOR[index])
          self.pos = None
          self.dir = None
          self.index = index
          self.view_size = 7
          self.carrying = None
          self.terminated = False
          self.started = True
          self.paused = False
  
      def render(self, img):
          c = COLORS[self.color]
          tri_fn = point_in_triangle(
              (0.12, 0.19),
              (0.87, 0.50),
              (0.12, 0.81),
          )
          # Rotate the agent based on its direction
          tri_fn = rotate_fn(tri_fn, cx=0.5, cy=0.5, theta=0.5 * math.pi * self.dir)
          fill_coords(img, tri_fn, c)
  
      def encode(self, current_agent = False):
          """Encode the a description of this object as a 3-tuple of integers"""
          if self.carrying:
              if current_agent:
                  return (OBJECT_TO_IDX[self.type], COLOR_TO_IDX[self.color], OBJECT_TO_IDX[self.carrying.type], COLOR_TO_IDX[self.carrying.color], self.dir, 1)
              else:
                  return (OBJECT_TO_IDX[self.type], COLOR_TO_IDX[self.color], OBJECT_TO_IDX[self.carrying.type], COLOR_TO_IDX[self.carrying.color], self.dir, 0)
  
          else:
              if current_agent:
                  return (OBJECT_TO_IDX[self.type], COLOR_TO_IDX[self.color], 0, 0, self.dir, 1)
              else:
                  return (OBJECT_TO_IDX[self.type], COLOR_TO_IDX[self.color], 0, 0, self.dir, 0)
  
  
      @property
      def dir_vec(self):
          """
          Get the direction vector for the agent, pointing in the direction
          of forward movement.
          """
  
          assert self.dir >= 0 and self.dir < 4
          return DIR_TO_VEC[self.dir]
  
      @property
      def right_vec(self):
          """
          Get the vector pointing to the right of the agent.
          """
  
          dx, dy = self.dir_vec
          return np.array((-dy, dx))
  
      @property
      def front_pos(self):
          """
          Get the position of the cell that is right in front of the agent
          """
  
          return self.pos + self.dir_vec
  
      def get_view_coords(self, i, j):
          """
          Translate and rotate absolute grid coordinates (i, j) into the
          agent's partially observable view (sub-grid). Note that the resulting
          coordinates may be negative or outside of the agent's view size.
          """
  
          ax, ay = self.pos
          dx, dy = self.dir_vec
          rx, ry = self.right_vec
  
          # Compute the absolute coordinates of the top-left view corner
          sz = self.view_size
          hs = self.view_size // 2
          tx = ax + (dx * (sz - 1)) - (rx * hs)
          ty = ay + (dy * (sz - 1)) - (ry * hs)
  
          lx = i - tx
          ly = j - ty
  
          # Project the coordinates of the object relative to the top-left
          # corner onto the agent's own coordinate system
          vx = (rx * lx + ry * ly)
          vy = -(dx * lx + dy * ly)
  
          return vx, vy
  
      def get_view_exts(self):
          """
          Get the extents of the square set of tiles visible to the agent
          Note: the bottom extent indices are not included in the set
          """
  
          # Facing right
          if self.dir == 0:
              topX = self.pos[0]
              topY = self.pos[1] - self.view_size // 2
          # Facing down
          elif self.dir == 1:
              topX = self.pos[0] - self.view_size // 2
              topY = self.pos[1]
          # Facing left
          elif self.dir == 2:
              topX = self.pos[0] - self.view_size + 1
              topY = self.pos[1] - self.view_size // 2
          # Facing up
          elif self.dir == 3:
              topX = self.pos[0] - self.view_size // 2
              topY = self.pos[1] - self.view_size + 1
          else:
              assert False, "invalid agent direction"
  
          botX = topX + self.view_size
          botY = topY + self.view_size
  
          return (topX, topY, botX, botY)
  
      def relative_coords(self, x, y):
          """
          Check if a grid position belongs to the agent's field of view, and returns the corresponding coordinates
          """
  
          vx, vy = self.get_view_coords(x, y)
  
          if vx < 0 or vy < 0 or vx >= self.view_size or vy >= self.view_size:
              return None
  
          return vx, vy
  
      def in_view(self, x, y):
          """
          check if a grid position is visible to the agent
          """
  
          return self.relative_coords(x, y) is not None
  
  
  class CollectGameEnv(MultiGridEnv):
      """
      Environment in which the agents have to collect the balls
      """
  
      def __init__(
          self,
          size=10,
          width=None,
          height=None,
          num_balls=[],
          agents_index = [],
          balls_index=[],
          balls_reward=[],
          zero_sum = False
  
      ):
          self.num_balls = num_balls
          self.balls_index = balls_index
          self.balls_reward = balls_reward
          self.zero_sum = zero_sum
  
          agents = []
          for i in agents_index:
              agents.append(Agent(i))
  
          super().__init__(
              grid_size=size,
              width=width,
              height=height,
              max_steps= 10000,
              # Set this to True for maximum speed
              see_through_walls=True,
              agents=agents
          )
  
      def _gen_grid(self, width, height):
          self.grid = Grid(width, height)
  
          # Generate the surrounding walls
          self.grid.horz_wall(0, 0)
          self.grid.horz_wall(0, height-1)
          self.grid.vert_wall(0, 0)
          self.grid.vert_wall(width-1, 0)
  
          for number, index, reward in zip(self.num_balls, self.balls_index, self.balls_reward):
              for i in range(number):
                  self.place_obj(Ball(index, reward))
  
          # Randomize the player start position and orientation
          for a in self.agents:
              self.place_agent(a)
  
      def _reward(self, i, rewards, reward=1):
          """
          Compute the reward to be given upon success
          """
          for j,a in enumerate(self.agents):
              if a.index==i or a.index==0:
                  rewards[j]+=reward
              if self.zero_sum:
                  if a.index!=i or a.index==0:
                      rewards[j] -= reward
  
      def _handle_pickup(self, i, rewards, fwd_pos, fwd_cell):
          if fwd_cell:
              if fwd_cell.can_pickup():
                  if fwd_cell.index in [0, self.agents[i].index]:
                      fwd_cell.cur_pos = np.array([-1, -1])
                      self.grid.set(*fwd_pos, None)
                      self._reward(i, rewards, fwd_cell.reward)
  
      def _handle_drop(self, i, rewards, fwd_pos, fwd_cell):
          pass
  
      def step(self, actions):
          obs, rewards, done, info = MiniGridEnv.step(self, actions)
          return obs, rewards, done, info
  
  
  class CollectGame4HEnv10x10N2(CollectGameEnv):
      def __init__(self):
          super().__init__(size=10,
          num_balls=[5],
          agents_index = [1,2,3],
          balls_index=[0],
          balls_reward=[1],
          zero_sum=True)
  
  
  register(
      id='MiniGrid-CollectGame4HEnv-10x10-N2-v0',
      entry_point='gym_minigrid.envs:CollectGame4HEnv10x10N2'
  )
  