From 44fdd9d207c649075819b512aa1ed8fb42aed013 Mon Sep 17 00:00:00 2001 From: Alexander Soare Date: Tue, 2 Jul 2024 11:49:45 +0100 Subject: [PATCH 01/10] add keypoints mode --- README.md | 7 +++++-- gym_pusht/envs/pusht.py | 44 +++++++++++++++++++++++++++++++++++++---- 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 04655fb..9cf5b33 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,9 @@ If `obs_type` is set to `state`, the observation space is a 5-dimensional vector environment: [agent_x, agent_y, block_x, block_y, block_angle]. The values are in the range [0, 512] for the agent and block positions and [0, 2*pi] for the block angle. +If `obs_type` is set to `keypoints` the observation space is a 16-dimensional vector representing the keypoint +locations of the T (in [x0, y0, x1, y1, ...] format). The values are in the range [0, 512]. + If `obs_type` is set to `pixels`, the observation space is a 96x96 RGB image of the environment. ### Rewards @@ -84,7 +87,7 @@ The episode terminates when the block is at least 95% in the goal zone. >>>> ``` -* `obs_type`: (str) The observation type. Can be either `state`, `pixels` or `pixels_agent_pos`. Default is `state`. +* `obs_type`: (str) The observation type. Can be either `state`, `kepoints`, `pixels` or `pixels_agent_pos`. Default is `state`. * `block_cog`: (tuple) The center of gravity of the block if different from the center of mass. Default is `None`. @@ -105,7 +108,7 @@ The episode terminates when the block is at least 95% in the goal zone. Passing the option `options["reset_to_state"]` will reset the environment to a specific state. > [!WARNING] -> For legacy compatibility, the inner fonctionning has been preserved, and the state set is not the same as the +> For legacy compatibility, the inner functioning has been preserved, and the state set is not the same as the > the one passed in the argument. ```python diff --git a/gym_pusht/envs/pusht.py b/gym_pusht/envs/pusht.py index 6c3b820..a641bcd 100644 --- a/gym_pusht/envs/pusht.py +++ b/gym_pusht/envs/pusht.py @@ -56,6 +56,10 @@ class PushTEnv(gym.Env): environment: [agent_x, agent_y, block_x, block_y, block_angle]. The values are in the range [0, 512] for the agent and block positions and [0, 2*pi] for the block angle. + If `obs_type` is set to `keypoints` the observation space is a 16-dimensional vector representing the keypoint + locations of the T (in [x0, y0, x1, y1, ...] format). The values are in the range [0, 512]. See `get_keypoints` for + a diagram showing the location of the keypoint indices. + If `obs_type` is set to `pixels`, the observation space is a 96x96 RGB image of the environment. ## Rewards @@ -84,7 +88,8 @@ class PushTEnv(gym.Env): >>>> ``` - * `obs_type`: (str) The observation type. Can be either `state`, `pixels` or `pixels_agent_pos`. Default is `state`. + * `obs_type`: (str) The observation type. Can be either `state`, `keypoints`, `pixels` or `pixels_agent_pos`. + Default is `state`. * `block_cog`: (tuple) The center of gravity of the block if different from the center of mass. Default is `None`. @@ -181,6 +186,12 @@ def _initialize_observation_space(self): high=np.array([512, 512, 512, 512, 2 * np.pi]), dtype=np.float64, ) + elif self.obs_type == "keypoints": + self.observation_space = spaces.Box( + low=np.zeros(8), + high=np.full((8,), 512), + dtype=np.float64, + ) elif self.obs_type == "pixels": self.observation_space = spaces.Box( low=0, high=255, shape=(self.observation_height, self.observation_width, 3), dtype=np.uint8 @@ -364,6 +375,9 @@ def get_obs(self): block_angle = self.block.angle % (2 * np.pi) return np.concatenate([agent_position, block_position, [block_angle]], dtype=np.float64) + if self.obs_type == "keypoints": + return self.get_keypoints().flatten() + pixels = self._render() if self.obs_type == "pixels": return pixels @@ -429,7 +443,7 @@ def _setup(self): def _set_state(self, state): self.agent.position = list(state[:2]) # Setting angle rotates with respect to center of mass, therefore will modify the geometric position if not - # the same as CoM. Therefore should theoritically set the angle first. But for compatibility with legacy data, + # the same as CoM. Therefore should theoretically set the angle first. But for compatibility with legacy data, # we do the opposite. self.block.position = list(state[2:4]) self.block.angle = state[4] @@ -454,8 +468,7 @@ def add_circle(space, position, radius): space.add(body, shape) return body - @staticmethod - def add_tee(space, position, angle, scale=30, color="LightSlateGray", mask=None): + def add_tee(self, space, position, angle, scale=30, color="LightSlateGray", mask=None): if mask is None: mask = pymunk.ShapeFilter.ALL_MASKS() mass = 1 @@ -486,4 +499,27 @@ def add_tee(space, position, angle, scale=30, color="LightSlateGray", mask=None) body.angle = angle body.friction = 1 space.add(body, shape1, shape2) + self._block_shapes = [shape1, shape2] return body + + def get_keypoints(self): + """Get a (8, 2) numpy array with the T keypoints. + + The T is composed of two rectangles each with 4 keypoints. + + 0───────────1 + │ │ + 3───4───5───2 + │ │ + │ │ + │ │ + │ │ + 7───6 + """ + keypoints = [] + for shape in self._block_shapes: + for v in shape.get_vertices(): + v = v.rotated(shape.body.angle) + v = v + shape.body.position + keypoints.append(np.array(v)) + return np.row_stack(keypoints) From 790b01694f7063f067659d7b2ee74d9a5b729b70 Mon Sep 17 00:00:00 2001 From: Alexander Soare Date: Tue, 2 Jul 2024 12:05:52 +0100 Subject: [PATCH 02/10] Revert `add_tee` to staticmethod --- gym_pusht/envs/pusht.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/gym_pusht/envs/pusht.py b/gym_pusht/envs/pusht.py index a641bcd..7ce146e 100644 --- a/gym_pusht/envs/pusht.py +++ b/gym_pusht/envs/pusht.py @@ -430,7 +430,7 @@ def _setup(self): # Add agent, block, and goal zone self.agent = self.add_circle(self.space, (256, 400), 15) - self.block = self.add_tee(self.space, (256, 300), 0) + self.block, self._block_shapes = self.add_tee(self.space, (256, 300), 0) self.goal_pose = np.array([256, 256, np.pi / 4]) # x, y, theta (in radians) if self.block_cog is not None: self.block.center_of_gravity = self.block_cog @@ -468,6 +468,7 @@ def add_circle(space, position, radius): space.add(body, shape) return body + @staticmethod def add_tee(self, space, position, angle, scale=30, color="LightSlateGray", mask=None): if mask is None: mask = pymunk.ShapeFilter.ALL_MASKS() @@ -499,8 +500,7 @@ def add_tee(self, space, position, angle, scale=30, color="LightSlateGray", mask body.angle = angle body.friction = 1 space.add(body, shape1, shape2) - self._block_shapes = [shape1, shape2] - return body + return body, [shape1, shape2] def get_keypoints(self): """Get a (8, 2) numpy array with the T keypoints. From 97e32f9d32db62c914b208d347bc5f4c26e0734b Mon Sep 17 00:00:00 2001 From: Alexander Soare Date: Tue, 2 Jul 2024 12:08:47 +0100 Subject: [PATCH 03/10] Make get_keypoints static --- gym_pusht/envs/pusht.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/gym_pusht/envs/pusht.py b/gym_pusht/envs/pusht.py index 7ce146e..7c5841c 100644 --- a/gym_pusht/envs/pusht.py +++ b/gym_pusht/envs/pusht.py @@ -376,7 +376,7 @@ def get_obs(self): return np.concatenate([agent_position, block_position, [block_angle]], dtype=np.float64) if self.obs_type == "keypoints": - return self.get_keypoints().flatten() + return self.get_keypoints(self._block_shapes).flatten() pixels = self._render() if self.obs_type == "pixels": @@ -469,7 +469,7 @@ def add_circle(space, position, radius): return body @staticmethod - def add_tee(self, space, position, angle, scale=30, color="LightSlateGray", mask=None): + def add_tee(space, position, angle, scale=30, color="LightSlateGray", mask=None): if mask is None: mask = pymunk.ShapeFilter.ALL_MASKS() mass = 1 @@ -502,7 +502,8 @@ def add_tee(self, space, position, angle, scale=30, color="LightSlateGray", mask space.add(body, shape1, shape2) return body, [shape1, shape2] - def get_keypoints(self): + @staticmethod + def get_keypoints(block_shapes): """Get a (8, 2) numpy array with the T keypoints. The T is composed of two rectangles each with 4 keypoints. @@ -517,7 +518,7 @@ def get_keypoints(self): 7───6 """ keypoints = [] - for shape in self._block_shapes: + for shape in block_shapes: for v in shape.get_vertices(): v = v.rotated(shape.body.angle) v = v + shape.body.position From 06f6fef9152fd1d43ec28ebebd8e2dc57478874a Mon Sep 17 00:00:00 2001 From: Alexander Soare Date: Tue, 2 Jul 2024 13:43:39 +0100 Subject: [PATCH 04/10] change keypoints -> environment_state and remove debug indexing --- README.md | 4 ++-- gym_pusht/envs/pusht.py | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 9cf5b33..82f26c6 100644 --- a/README.md +++ b/README.md @@ -56,7 +56,7 @@ If `obs_type` is set to `state`, the observation space is a 5-dimensional vector environment: [agent_x, agent_y, block_x, block_y, block_angle]. The values are in the range [0, 512] for the agent and block positions and [0, 2*pi] for the block angle. -If `obs_type` is set to `keypoints` the observation space is a 16-dimensional vector representing the keypoint +If `obs_type` is set to `environment_state` the observation space is a 16-dimensional vector representing the keypoint locations of the T (in [x0, y0, x1, y1, ...] format). The values are in the range [0, 512]. If `obs_type` is set to `pixels`, the observation space is a 96x96 RGB image of the environment. @@ -87,7 +87,7 @@ The episode terminates when the block is at least 95% in the goal zone. >>>> ``` -* `obs_type`: (str) The observation type. Can be either `state`, `kepoints`, `pixels` or `pixels_agent_pos`. Default is `state`. +* `obs_type`: (str) The observation type. Can be either `state`, `environment_state`, `pixels` or `pixels_agent_pos`. Default is `state`. * `block_cog`: (tuple) The center of gravity of the block if different from the center of mass. Default is `None`. diff --git a/gym_pusht/envs/pusht.py b/gym_pusht/envs/pusht.py index 7c5841c..4449bb2 100644 --- a/gym_pusht/envs/pusht.py +++ b/gym_pusht/envs/pusht.py @@ -56,9 +56,9 @@ class PushTEnv(gym.Env): environment: [agent_x, agent_y, block_x, block_y, block_angle]. The values are in the range [0, 512] for the agent and block positions and [0, 2*pi] for the block angle. - If `obs_type` is set to `keypoints` the observation space is a 16-dimensional vector representing the keypoint - locations of the T (in [x0, y0, x1, y1, ...] format). The values are in the range [0, 512]. See `get_keypoints` for - a diagram showing the location of the keypoint indices. + If `obs_type` is set to `environment_state` the observation space is a 16-dimensional vector representing the + keypoint locations of the T (in [x0, y0, x1, y1, ...] format). The values are in the range [0, 512]. See + `get_keypoints` for a diagram showing the location of the keypoint indices. If `obs_type` is set to `pixels`, the observation space is a 96x96 RGB image of the environment. @@ -186,7 +186,7 @@ def _initialize_observation_space(self): high=np.array([512, 512, 512, 512, 2 * np.pi]), dtype=np.float64, ) - elif self.obs_type == "keypoints": + elif self.obs_type == "environment_state": self.observation_space = spaces.Box( low=np.zeros(8), high=np.full((8,), 512), @@ -375,7 +375,7 @@ def get_obs(self): block_angle = self.block.angle % (2 * np.pi) return np.concatenate([agent_position, block_position, [block_angle]], dtype=np.float64) - if self.obs_type == "keypoints": + if self.obs_type == "environment_state": return self.get_keypoints(self._block_shapes).flatten() pixels = self._render() From 6b16c3d8830be71f2c00f25ce2549eb67e453b80 Mon Sep 17 00:00:00 2001 From: Alexander Soare Date: Tue, 2 Jul 2024 14:21:29 +0100 Subject: [PATCH 05/10] fix observation space --- gym_pusht/envs/pusht.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gym_pusht/envs/pusht.py b/gym_pusht/envs/pusht.py index 4449bb2..e88def0 100644 --- a/gym_pusht/envs/pusht.py +++ b/gym_pusht/envs/pusht.py @@ -188,8 +188,8 @@ def _initialize_observation_space(self): ) elif self.obs_type == "environment_state": self.observation_space = spaces.Box( - low=np.zeros(8), - high=np.full((8,), 512), + low=np.zeros(16), + high=np.full((16,), 512), dtype=np.float64, ) elif self.obs_type == "pixels": From aaabdde8b85984cfb82ca6445c66508d65312656 Mon Sep 17 00:00:00 2001 From: Alexander Soare Date: Tue, 2 Jul 2024 14:32:11 +0100 Subject: [PATCH 06/10] use environment_state_agent_pos --- README.md | 6 ++++-- gym_pusht/envs/pusht.py | 34 ++++++++++++++++++++++++---------- 2 files changed, 28 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 82f26c6..c5ed48e 100644 --- a/README.md +++ b/README.md @@ -56,8 +56,10 @@ If `obs_type` is set to `state`, the observation space is a 5-dimensional vector environment: [agent_x, agent_y, block_x, block_y, block_angle]. The values are in the range [0, 512] for the agent and block positions and [0, 2*pi] for the block angle. -If `obs_type` is set to `environment_state` the observation space is a 16-dimensional vector representing the keypoint -locations of the T (in [x0, y0, x1, y1, ...] format). The values are in the range [0, 512]. +If `obs_type` is set to `environment_state_agent_pos` the observation space is a dictionary with: + - `environment_state`: 16-dimensional vector representing the keypoint locations of the T (in [x0, y0, x1, y1, ...] + format). The values are in the range [0, 512]. + - `agent_pos`: A 2-dimensional vector representing the position of the robot end-effector. If `obs_type` is set to `pixels`, the observation space is a 96x96 RGB image of the environment. diff --git a/gym_pusht/envs/pusht.py b/gym_pusht/envs/pusht.py index e88def0..692ea4f 100644 --- a/gym_pusht/envs/pusht.py +++ b/gym_pusht/envs/pusht.py @@ -56,9 +56,11 @@ class PushTEnv(gym.Env): environment: [agent_x, agent_y, block_x, block_y, block_angle]. The values are in the range [0, 512] for the agent and block positions and [0, 2*pi] for the block angle. - If `obs_type` is set to `environment_state` the observation space is a 16-dimensional vector representing the - keypoint locations of the T (in [x0, y0, x1, y1, ...] format). The values are in the range [0, 512]. See - `get_keypoints` for a diagram showing the location of the keypoint indices. + If `obs_type` is set to `environment_state_agent_pos` the observation space is a dictionary with: + - `environment_state`: 16-dimensional vector representing the keypoint locations of the T (in [x0, y0, x1, y1, ...] + format). The values are in the range [0, 512]. See `get_keypoints` for a diagram showing the location of the + keypoint indices. + - `agent_pos`: A 2-dimensional vector representing the position of the robot end-effector. If `obs_type` is set to `pixels`, the observation space is a 96x96 RGB image of the environment. @@ -186,11 +188,20 @@ def _initialize_observation_space(self): high=np.array([512, 512, 512, 512, 2 * np.pi]), dtype=np.float64, ) - elif self.obs_type == "environment_state": - self.observation_space = spaces.Box( - low=np.zeros(16), - high=np.full((16,), 512), - dtype=np.float64, + elif self.obs_type == "environment_state_agent_pos": + self.observation_space = spaces.Dict( + { + "environment_state": spaces.Box( + low=np.zeros(16), + high=np.full((16,), 512), + dtype=np.float64, + ), + "agent_pos": spaces.Box( + low=np.array([0, 0]), + high=np.array([512, 512]), + dtype=np.float64, + ), + }, ) elif self.obs_type == "pixels": self.observation_space = spaces.Box( @@ -375,8 +386,11 @@ def get_obs(self): block_angle = self.block.angle % (2 * np.pi) return np.concatenate([agent_position, block_position, [block_angle]], dtype=np.float64) - if self.obs_type == "environment_state": - return self.get_keypoints(self._block_shapes).flatten() + if self.obs_type == "environment_state_agent_pos": + return { + "environment_state": self.get_keypoints(self._block_shapes).flatten(), + "agent_pos": np.array(self.agent.position), + } pixels = self._render() if self.obs_type == "pixels": From 312c8b391a7f33014595cfece7867543c37a055f Mon Sep 17 00:00:00 2001 From: Alexander Soare Date: Tue, 2 Jul 2024 14:50:55 +0100 Subject: [PATCH 07/10] update error message --- gym_pusht/envs/pusht.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/gym_pusht/envs/pusht.py b/gym_pusht/envs/pusht.py index 692ea4f..6fc13e5 100644 --- a/gym_pusht/envs/pusht.py +++ b/gym_pusht/envs/pusht.py @@ -225,7 +225,8 @@ def _initialize_observation_space(self): ) else: raise ValueError( - f"Unknown obs_type {self.obs_type}. Must be one of [pixels, state, pixels_agent_pos]" + f"Unknown obs_type {self.obs_type}. Must be one of [pixels, state, environment_state_agent_pos, " + "pixels_agent_pos]" ) def _get_coverage(self): From 291cca77e8cca3f148628e3420f308609904a1c4 Mon Sep 17 00:00:00 2001 From: Alexander Soare Date: Wed, 3 Jul 2024 14:46:55 +0100 Subject: [PATCH 08/10] fix add tee --- gym_pusht/envs/pusht.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gym_pusht/envs/pusht.py b/gym_pusht/envs/pusht.py index 6fc13e5..73c9a35 100644 --- a/gym_pusht/envs/pusht.py +++ b/gym_pusht/envs/pusht.py @@ -511,8 +511,8 @@ def add_tee(space, position, angle, scale=30, color="LightSlateGray", mask=None) shape1.filter = pymunk.ShapeFilter(mask=mask) shape2.filter = pymunk.ShapeFilter(mask=mask) body.center_of_gravity = (shape1.center_of_gravity + shape2.center_of_gravity) / 2 - body.position = position body.angle = angle + body.position = position body.friction = 1 space.add(body, shape1, shape2) return body, [shape1, shape2] From 9d19f907f2c63328df75850920982be1b0cf7e95 Mon Sep 17 00:00:00 2001 From: Alexander Soare Date: Thu, 4 Jul 2024 13:33:28 +0100 Subject: [PATCH 09/10] fix variable name in documentation --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index c5ed48e..74eadd9 100644 --- a/README.md +++ b/README.md @@ -89,7 +89,7 @@ The episode terminates when the block is at least 95% in the goal zone. >>>> ``` -* `obs_type`: (str) The observation type. Can be either `state`, `environment_state`, `pixels` or `pixels_agent_pos`. Default is `state`. +* `obs_type`: (str) The observation type. Can be either `state`, `environment_state_agent_pos`, `pixels` or `pixels_agent_pos`. Default is `state`. * `block_cog`: (tuple) The center of gravity of the block if different from the center of mass. Default is `None`. From c6ee94cfbcab19e82863955bda778699a22cd37f Mon Sep 17 00:00:00 2001 From: Alexander Soare Date: Thu, 4 Jul 2024 16:02:56 +0100 Subject: [PATCH 10/10] add test --- tests/test_env.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_env.py b/tests/test_env.py index f747e2a..91562ba 100644 --- a/tests/test_env.py +++ b/tests/test_env.py @@ -11,6 +11,7 @@ ("PushT-v0", "state"), ("PushT-v0", "pixels"), ("PushT-v0", "pixels_agent_pos"), + ("PushT-v0", "environment_state_agent_pos"), ], ) def test_env(env_task, obs_type):