Added radius rescaling to simple env (#1213)

pettingzoo/mpe/_mpe_utils/simple_env.py

@@ -42,6 +42,7 @@ def __init__(
         render_mode=None,
         continuous_actions=False,
         local_ratio=None,
+        dynamic_rescaling=False,
     ):
         super().__init__()
 
@@ -66,6 +67,7 @@ def __init__(
         self.world = world
         self.continuous_actions = continuous_actions
         self.local_ratio = local_ratio
+        self.dynamic_rescaling = dynamic_rescaling
 
         self.scenario.reset_world(self.world, self.np_random)
 
@@ -116,6 +118,11 @@ def __init__(
             dtype=np.float32,
         )
 
+        # Get the original cam_range
+        # This will be used to scale the rendering
+        all_poses = [entity.state.p_pos for entity in self.world.entities]
+        self.original_cam_range = np.max(np.abs(np.array(all_poses)))
+
         self.steps = 0
 
         self.current_actions = [None] * self.num_agents
@@ -295,6 +302,10 @@ def draw(self):
         all_poses = [entity.state.p_pos for entity in self.world.entities]
         cam_range = np.max(np.abs(np.array(all_poses)))
 
+        # The scaling factor is used for dynamic rescaling of the rendering - a.k.a Zoom In/Zoom Out effect
+        # The 0.9 is a factor to keep the entities from appearing "too" out-of-bounds
+        scaling_factor = 0.9 * self.original_cam_range / cam_range
+
         # update geometry and text positions
         text_line = 0
         for e, entity in enumerate(self.world.entities):
@@ -309,12 +320,15 @@ def draw(self):
             y = (y / cam_range) * self.height // 2 * 0.9
             x += self.width // 2
             y += self.height // 2
-            pygame.draw.circle(
-                self.screen, entity.color * 200, (x, y), entity.size * 350
-            )  # 350 is an arbitrary scale factor to get pygame to render similar sizes as pyglet
-            pygame.draw.circle(
-                self.screen, (0, 0, 0), (x, y), entity.size * 350, 1
-            )  # borders
+
+            # 350 is an arbitrary scale factor to get pygame to render similar sizes as pyglet
+            if self.dynamic_rescaling:
+                radius = entity.size * 350 * scaling_factor
+            else:
+                radius = entity.size * 350
+
+            pygame.draw.circle(self.screen, entity.color * 200, (x, y), radius)
+            pygame.draw.circle(self.screen, (0, 0, 0), (x, y), radius, 1)  # borders
             assert (
                 0 < x < self.width and 0 < y < self.height
             ), f"Coordinates {(x, y)} are out of bounds."

pettingzoo/mpe/simple/simple.py

@@ -31,7 +31,7 @@
 ### Arguments
 
 ``` python
-simple_v3.env(max_cycles=25, continuous_actions=False)
+simple_v3.env(max_cycles=25, continuous_actions=False, dynamic_rescaling=False)
 ```
 
 
@@ -40,6 +40,8 @@
 
 `continuous_actions`: Whether agent action spaces are discrete(default) or continuous
 
+`dynamic_rescaling`: Whether to rescale the size of agents and landmarks based on the screen size
+
 """
 
 import numpy as np
@@ -52,7 +54,13 @@
 
 
 class raw_env(SimpleEnv, EzPickle):
-    def __init__(self, max_cycles=25, continuous_actions=False, render_mode=None):
+    def __init__(
+        self,
+        max_cycles=25,
+        continuous_actions=False,
+        render_mode=None,
+        dynamic_rescaling=False,
+    ):
         EzPickle.__init__(
             self,
             max_cycles=max_cycles,
@@ -68,6 +76,7 @@ def __init__(self, max_cycles=25, continuous_actions=False, render_mode=None):
             render_mode=render_mode,
             max_cycles=max_cycles,
             continuous_actions=continuous_actions,
+            dynamic_rescaling=dynamic_rescaling,
         )
         self.metadata["name"] = "simple_v3"
 

pettingzoo/mpe/simple_adversary/simple_adversary.py

@@ -39,7 +39,7 @@
 ### Arguments
 
 ``` python
-simple_adversary_v3.env(N=2, max_cycles=25, continuous_actions=False)
+simple_adversary_v3.env(N=2, max_cycles=25, continuous_actions=False, dynamic_rescaling=False)
 ```
 
 
@@ -50,6 +50,8 @@
 
 `continuous_actions`: Whether agent action spaces are discrete(default) or continuous
 
+`dynamic_rescaling`: Whether to rescale the size of agents and landmarks based on the screen size
+
 """
 
 import numpy as np
@@ -62,7 +64,14 @@
 
 
 class raw_env(SimpleEnv, EzPickle):
-    def __init__(self, N=2, max_cycles=25, continuous_actions=False, render_mode=None):
+    def __init__(
+        self,
+        N=2,
+        max_cycles=25,
+        continuous_actions=False,
+        render_mode=None,
+        dynamic_rescaling=False,
+    ):
         EzPickle.__init__(
             self,
             N=N,
@@ -79,6 +88,7 @@ def __init__(self, N=2, max_cycles=25, continuous_actions=False, render_mode=Non
             render_mode=render_mode,
             max_cycles=max_cycles,
             continuous_actions=continuous_actions,
+            dynamic_rescaling=dynamic_rescaling,
         )
         self.metadata["name"] = "simple_adversary_v3"
 

pettingzoo/mpe/simple_crypto/simple_crypto.py

@@ -45,7 +45,7 @@
 ### Arguments
 
 ``` python
-simple_crypto_v3.env(max_cycles=25, continuous_actions=False)
+simple_crypto_v3.env(max_cycles=25, continuous_actions=False, dynamic_rescaling=False)
 ```
 
 
@@ -54,6 +54,8 @@
 
 `continuous_actions`: Whether agent action spaces are discrete(default) or continuous
 
+`dynamic_rescaling`: Whether to rescale the size of agents and landmarks based on the screen size
+
 """
 
 import numpy as np
@@ -73,7 +75,13 @@
 
 
 class raw_env(SimpleEnv, EzPickle):
-    def __init__(self, max_cycles=25, continuous_actions=False, render_mode=None):
+    def __init__(
+        self,
+        max_cycles=25,
+        continuous_actions=False,
+        render_mode=None,
+        dynamic_rescaling=False,
+    ):
         EzPickle.__init__(
             self,
             max_cycles=max_cycles,
@@ -89,6 +97,7 @@ def __init__(self, max_cycles=25, continuous_actions=False, render_mode=None):
             render_mode=render_mode,
             max_cycles=max_cycles,
             continuous_actions=continuous_actions,
+            dynamic_rescaling=dynamic_rescaling,
         )
         self.metadata["name"] = "simple_crypto_v3"
 

pettingzoo/mpe/simple_push/simple_push.py

@@ -38,13 +38,16 @@
 ### Arguments
 
 ``` python
-simple_push_v3.env(max_cycles=25, continuous_actions=False)
+simple_push_v3.env(max_cycles=25, continuous_actions=False, dynamic_rescaling=False)
 ```
 
 
 
 `max_cycles`:  number of frames (a step for each agent) until game terminates
 
+`dynamic_rescaling`: Whether to rescale the size of agents and landmarks based on the screen size
+
+
 """
 
 import numpy as np
@@ -57,7 +60,13 @@
 
 
 class raw_env(SimpleEnv, EzPickle):
-    def __init__(self, max_cycles=25, continuous_actions=False, render_mode=None):
+    def __init__(
+        self,
+        max_cycles=25,
+        continuous_actions=False,
+        render_mode=None,
+        dynamic_rescaling=False,
+    ):
         EzPickle.__init__(
             self,
             max_cycles=max_cycles,
@@ -73,6 +82,7 @@ def __init__(self, max_cycles=25, continuous_actions=False, render_mode=None):
             render_mode=render_mode,
             max_cycles=max_cycles,
             continuous_actions=continuous_actions,
+            dynamic_rescaling=dynamic_rescaling,
         )
         self.metadata["name"] = "simple_push_v3"
 

pettingzoo/mpe/simple_reference/simple_reference.py

@@ -40,7 +40,7 @@
 
 
 ``` python
-simple_reference_v3.env(local_ratio=0.5, max_cycles=25, continuous_actions=False)
+simple_reference_v3.env(local_ratio=0.5, max_cycles=25, continuous_actions=False, dynamic_rescaling=False)
 ```
 
 
@@ -51,6 +51,8 @@
 
 `continuous_actions`: Whether agent action spaces are discrete(default) or continuous
 
+`dynamic_rescaling`: Whether to rescale the size of agents and landmarks based on the screen size
+
 """
 
 import numpy as np
@@ -64,7 +66,12 @@
 
 class raw_env(SimpleEnv, EzPickle):
     def __init__(
-        self, local_ratio=0.5, max_cycles=25, continuous_actions=False, render_mode=None
+        self,
+        local_ratio=0.5,
+        max_cycles=25,
+        continuous_actions=False,
+        render_mode=None,
+        dynamic_rescaling=False,
     ):
         EzPickle.__init__(
             self,
@@ -86,6 +93,7 @@ def __init__(
             max_cycles=max_cycles,
             continuous_actions=continuous_actions,
             local_ratio=local_ratio,
+            dynamic_rescaling=dynamic_rescaling,
         )
         self.metadata["name"] = "simple_reference_v3"
 

pettingzoo/mpe/simple_speaker_listener/simple_speaker_listener.py

@@ -37,7 +37,7 @@
 ### Arguments
 
 ``` python
-simple_speaker_listener_v4.env(max_cycles=25, continuous_actions=False)
+simple_speaker_listener_v4.env(max_cycles=25, continuous_actions=False, dynamic_rescaling=False)
 ```
 
 
@@ -46,6 +46,8 @@
 
 `continuous_actions`: Whether agent action spaces are discrete(default) or continuous
 
+`dynamic_rescaling`: Whether to rescale the size of agents and landmarks based on the screen size
+
 """
 
 import numpy as np
@@ -58,7 +60,13 @@
 
 
 class raw_env(SimpleEnv, EzPickle):
-    def __init__(self, max_cycles=25, continuous_actions=False, render_mode=None):
+    def __init__(
+        self,
+        max_cycles=25,
+        continuous_actions=False,
+        render_mode=None,
+        dynamic_rescaling=False,
+    ):
         EzPickle.__init__(
             self,
             max_cycles=max_cycles,
@@ -74,6 +82,7 @@ def __init__(self, max_cycles=25, continuous_actions=False, render_mode=None):
             render_mode=render_mode,
             max_cycles=max_cycles,
             continuous_actions=continuous_actions,
+            dynamic_rescaling=dynamic_rescaling,
         )
         self.metadata["name"] = "simple_speaker_listener_v4"
 

pettingzoo/mpe/simple_spread/simple_spread.py

@@ -36,7 +36,7 @@
 ### Arguments
 
 ``` python
-simple_spread_v3.env(N=3, local_ratio=0.5, max_cycles=25, continuous_actions=False)
+simple_spread_v3.env(N=3, local_ratio=0.5, max_cycles=25, continuous_actions=False, dynamic_rescaling=False)
 ```
 
 
@@ -49,6 +49,8 @@
 
 `continuous_actions`: Whether agent action spaces are discrete(default) or continuous
 
+`dynamic_rescaling`: Whether to rescale the size of agents and landmarks based on the screen size
+
 """
 
 import numpy as np
@@ -68,6 +70,7 @@ def __init__(
         max_cycles=25,
         continuous_actions=False,
         render_mode=None,
+        dynamic_rescaling=False,
     ):
         EzPickle.__init__(
             self,
@@ -90,6 +93,7 @@ def __init__(
             max_cycles=max_cycles,
             continuous_actions=continuous_actions,
             local_ratio=local_ratio,
+            dynamic_rescaling=dynamic_rescaling,
         )
         self.metadata["name"] = "simple_spread_v3"
 

pettingzoo/mpe/simple_tag/simple_tag.py

@@ -45,7 +45,7 @@ def bound(x):
 ### Arguments
 
 ``` python
-simple_tag_v3.env(num_good=1, num_adversaries=3, num_obstacles=2, max_cycles=25, continuous_actions=False)
+simple_tag_v3.env(num_good=1, num_adversaries=3, num_obstacles=2, max_cycles=25, continuous_actions=False, dynamic_rescaling=False)
 ```
 
 
@@ -60,6 +60,8 @@ def bound(x):
 
 `continuous_actions`: Whether agent action spaces are discrete(default) or continuous
 
+`dynamic_rescaling`: Whether to rescale the size of agents and landmarks based on the screen size
+
 """
 
 import numpy as np
@@ -80,6 +82,7 @@ def __init__(
         max_cycles=25,
         continuous_actions=False,
         render_mode=None,
+        dynamic_rescaling=False,
     ):
         EzPickle.__init__(
             self,
@@ -99,6 +102,7 @@ def __init__(
             render_mode=render_mode,
             max_cycles=max_cycles,
             continuous_actions=continuous_actions,
+            dynamic_rescaling=dynamic_rescaling,
         )
         self.metadata["name"] = "simple_tag_v3"