From eef1c62b76cc802f7521f5b3dd1736229c894ecb Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:33:50 +0200 Subject: [PATCH 01/14] fix: repair broken gymnax.wrappers.purerl import, declare chex dependency Newer gymnax no longer re-exports functools.partial/typing.Optional/Tuple/ Union/chex/environment/spaces through gymnax.wrappers.purerl, so importing wrappers.py raised an ImportError on any current gymnax install. Import each symbol from its actual source instead. chex is used here (and in meta_cartpole.py) but was never listed in requirements.txt. --- popjym/wrappers.py | 16 ++++++---------- requirements/requirements.txt | 1 + 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/popjym/wrappers.py b/popjym/wrappers.py index 6c02c19..3451432 100644 --- a/popjym/wrappers.py +++ b/popjym/wrappers.py @@ -1,16 +1,12 @@ +from functools import partial +from typing import Optional, Tuple, Union + +import chex import jax import jax.numpy as jnp from flax import struct -from gymnax.wrappers.purerl import ( - GymnaxWrapper, - Optional, - Tuple, - Union, - chex, - environment, - partial, - spaces, -) +from gymnax.environments import environment, spaces +from gymnax.wrappers.purerl import GymnaxWrapper class AliasPrevAction(GymnaxWrapper): diff --git a/requirements/requirements.txt b/requirements/requirements.txt index aaf6f95..898b40f 100755 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -1,3 +1,4 @@ +chex flax gymnax jax From ed8c1c7677338ac995141e1c0d1116ffbde4093f Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:33:52 +0200 Subject: [PATCH 02/14] feat: auto-wrap envs that declare obs_requires_prev_action Several envs ported from popgym rely on the reference implementation's obs_requires_prev_action=True contract, which concatenates the agent's previous action onto the observation. That wrapping was never applied, silently changing the task. make() now checks for the attribute and auto-applies AliasPrevActionV2, printing a one-line notice so the wrapping is visible. No env declares the attribute yet; wired up per-env in the following commits. --- popjym/registration.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/popjym/registration.py b/popjym/registration.py index e6eebc5..9391241 100644 --- a/popjym/registration.py +++ b/popjym/registration.py @@ -127,6 +127,14 @@ def make(env_id: str, **env_kwargs): else: raise ValueError("Environment ID is not registered.") + if getattr(env, "obs_requires_prev_action", False): + from popjym.wrappers import AliasPrevActionV2 + print( + f"[popjym] {env_id} declares obs_requires_prev_action=True; " + f"auto-wrapping with AliasPrevActionV2 (observation includes prev action)." + ) + env = AliasPrevActionV2(env) + return env, env.default_params From fc0dc1626265246ca2629bae686f71363ad46252 Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:34:01 +0200 Subject: [PATCH 03/14] fix: Autoencode observation inversion, OOB read, off-by-one episode length - WATCH/PLAY observation was inverted: the env returned the card during PLAY and zeros during WATCH, exactly backwards from popgym's semantics (jnp.where(play, play_obs, watch_obs) now matches). - state.cards[state.timestep] could index out of bounds once state.timestep >= num_cards during PLAY; index by state.cards[state.timestep % num_cards] instead (value is unused once play_obs is selected, but the read itself must stay in-bounds). - Episode was one step too long, so the terminal step read a JAX-clamped out-of-bounds card. terminated now fires at state.timestep == num_cards * 2 - 2 and the PLAY-phase reward gate/index are shifted to match, giving 2*num_cards - 1 total steps as in popgym. --- popjym/environments/popgym_autoencode.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/popjym/environments/popgym_autoencode.py b/popjym/environments/popgym_autoencode.py index 8cfe02f..19e6abf 100644 --- a/popjym/environments/popgym_autoencode.py +++ b/popjym/environments/popgym_autoencode.py @@ -37,12 +37,11 @@ def step_env( reward_scale = 1.0 / (num_cards) - terminated = state.timestep == num_cards * 2 - play = state.timestep >= num_cards + terminated = state.timestep == num_cards * 2 - 2 + play = state.timestep >= num_cards - 1 reward = jnp.where( - # jnp.logical_and(play, state.cards[-(state.timestep - num_cards - 1)] == action), - jnp.flip(state.cards, axis=0)[state.timestep - num_cards] == action, + jnp.flip(state.cards, axis=0)[state.timestep - (num_cards - 1)] == action, reward_scale, -reward_scale, ) @@ -71,10 +70,11 @@ def reset_env(self, key: chex.PRNGKey, params: EnvParams) -> Tuple[chex.Array, E def get_obs(self, state: EnvState) -> chex.Array: """Returns observation from the state.""" - play = state.timestep >= self.decksize * self.num_decks + num_cards = self.decksize * self.num_decks + play = state.timestep >= num_cards play_obs = jnp.zeros((self.num_suits,)) - watch_obs = play_obs.at[state.cards[state.timestep]].set(1) - obs = jnp.where(play, watch_obs, play_obs) + watch_obs = play_obs.at[state.cards[state.timestep % num_cards]].set(1) + obs = jnp.where(play, play_obs, watch_obs) return obs def action_space(self, params: Optional[EnvParams] = None) -> spaces.Discrete: From a4dcd046d99406efc4b7750d54ada883b46f2ad6 Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:34:03 +0200 Subject: [PATCH 04/14] fix: Battleship ship placement bias, missing prev-action contract - Ship placement sampled directions non-uniformly, biasing ships toward the upper-left. is_valid_placement/place_ship_on_board now build all 4 directional deltas and sample a valid placement uniformly via jax.random.choice, matching popgym's placement distribution. - Declare obs_requires_prev_action = True so make() auto-wraps the env with AliasPrevActionV2, matching popgym's contract for this env. --- popjym/environments/popgym_battleship.py | 53 ++++++++++-------------- 1 file changed, 23 insertions(+), 30 deletions(-) diff --git a/popjym/environments/popgym_battleship.py b/popjym/environments/popgym_battleship.py index 4a47a9c..f6e725b 100644 --- a/popjym/environments/popgym_battleship.py +++ b/popjym/environments/popgym_battleship.py @@ -10,22 +10,19 @@ def is_valid_placement(board, row, col, direction, ship_size): """Check if a placement is valid without modifying the board.""" - board_shape = board.shape - - # Slice the board - horizontal_board = jax.lax.dynamic_slice(board, (row, col), (1, ship_size)) - vertical_board = jax.lax.dynamic_slice(board, (row, col), (ship_size, 1)) - - # Check validities - horizontal_validity = jnp.logical_and( - col + ship_size <= board_shape[1], jnp.all(horizontal_board == 0) + board_size = board.shape[0] + offsets = jnp.arange(ship_size) + deltas = jnp.array([[0, 1], [0, -1], [1, 0], [-1, 0]]) + drow, dcol = deltas[direction] + rows = row + drow * offsets + cols = col + dcol * offsets + in_bounds = jnp.all( + (rows >= 0) & (rows < board_size) & (cols >= 0) & (cols < board_size) ) - - vertical_validity = jnp.logical_and( - row + ship_size <= board_shape[0], jnp.all(vertical_board == 0) - ) - - return jnp.where(direction == 0, horizontal_validity, vertical_validity) + safe_rows = jnp.clip(rows, 0, board_size - 1) + safe_cols = jnp.clip(cols, 0, board_size - 1) + empty = jnp.all(board[safe_rows, safe_cols] == 0) + return jnp.logical_and(in_bounds, empty) vectorized_validity_check = jax.vmap( @@ -34,28 +31,23 @@ def is_valid_placement(board, row, col, direction, ship_size): in_axes=(None, None, 0, None, None), ), in_axes=(None, None, None, 0, None), -) # Why +) def place_ship_on_board(board, row, col, direction, ship_size): """Place a ship on the board at the given position and direction.""" - # Generate the horizontal and vertical ship placements - horizontal_ship = jnp.ones((1, ship_size)) - vertical_ship = jnp.ones((ship_size, 1)) - - # Create boards with the ship placed in each direction - horizontal_board = jax.lax.dynamic_update_slice(board, horizontal_ship, (row, col)) - vertical_board = jax.lax.dynamic_update_slice(board, vertical_ship, (row, col)) - - # Use `lax.select` to choose the appropriate board based on the direction - updated_board = jax.lax.select(direction == 0, horizontal_board, vertical_board) - - return updated_board + board_size = board.shape[0] + offsets = jnp.arange(ship_size) + deltas = jnp.array([[0, 1], [0, -1], [1, 0], [-1, 0]]) + drow, dcol = deltas[direction] + rows = jnp.clip(row + drow * offsets, 0, board_size - 1) + cols = jnp.clip(col + dcol * offsets, 0, board_size - 1) + return board.at[rows, cols].set(1.0) def place_random_ship_on_board(rng, board, ship_size): size = board.shape[0] - dirs = jnp.array([0, 1]) + dirs = jnp.arange(4) rows = jnp.arange(size) cols = jnp.arange(size) valid_spots = vectorized_validity_check(board, rows, cols, dirs, ship_size) @@ -68,7 +60,6 @@ def place_random_ship_on_board(rng, board, ship_size): (rand_valid % (size * size)) // size, (rand_valid % (size * size)) % size, ) - # print(is_valid_placement(board, row, col, direction, ship_size)) board = place_ship_on_board(board, row, col, direction, ship_size) return board @@ -95,6 +86,8 @@ class EnvParams: class Battleship(environment.Environment): + obs_requires_prev_action = True + def __init__(self, board_size=8): super().__init__() self.board_size = board_size From 962fd81070a14443b320575d29019ca229d23d8f Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:34:05 +0200 Subject: [PATCH 05/14] fix: Concentration episode_length dtype, missing prev-action contract - episode_length was a jnp float, making the timestep == episode_length termination check fragile; cast to a Python int. - Declare obs_requires_prev_action = True so make() auto-wraps the env with AliasPrevActionV2, matching popgym's contract for this env. - Drop a stray leftover comment. --- popjym/environments/popgym_concentration.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/popjym/environments/popgym_concentration.py b/popjym/environments/popgym_concentration.py index 8b22a94..90381c9 100644 --- a/popjym/environments/popgym_concentration.py +++ b/popjym/environments/popgym_concentration.py @@ -21,6 +21,8 @@ class EnvParams: class Concentration(environment.Environment): + obs_requires_prev_action = True + def __init__(self, num_decks=1, num_types=2): super().__init__() self.decksize = 52 @@ -28,8 +30,8 @@ def __init__(self, num_decks=1, num_types=2): self.num_decks = num_decks self.num_types = num_types self.num_cards = self.decksize * self.num_decks - self.episode_length = jnp.ceil( - 2 * self.num_cards - (self.num_cards / (2.0 * self.num_cards - 1)) + self.episode_length = int( + jnp.ceil(2 * self.num_cards - (self.num_cards / (2.0 * self.num_cards - 1))) ) self.success_reward_scale = 1.0 / (self.num_cards // 2) self.failure_reward_scale = -1.0 / (self.episode_length) @@ -55,7 +57,7 @@ def step_env( # IF TRYING CARD ALREADY UP reward = jnp.where( trying_card_already_up, - jnp.sum(new_in_play) * self.failure_reward_scale, # WHY IS IT SCALED BY NUM IN PLAY? + jnp.sum(new_in_play) * self.failure_reward_scale, 0.0, ) new_in_play = jnp.where(trying_card_already_up, jnp.zeros_like(new_in_play), new_in_play) From de7e75ac3e633eb33c85baefac6e3140851ec21a Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:34:08 +0200 Subject: [PATCH 06/14] fix: CountRecall hardcoded action space, timing, reset, reward scale, length - action_space was hardcoded Discrete(2) regardless of difficulty, making Medium/Hard unsolvable; now spaces.Discrete(self.max_num + 1) as popgym does. - prev_count must be read from running_count before the current card is tallied; it was being read after, off by one card from popgym's prev_query/prev_count-then-deal-then-count ordering. - Initial running_count didn't tally the first dealt card on reset; seed it via running_count.at[value_cards[0]].add(1) as popgym does in __init__. - Reward scale used the wrong denominator; corrected to 1.0 / (num_cards - 1) matching popgym's 1/max_episode_length. - Episode was one step too long; terminate at new_state.timestep == self.num_cards - 1. --- popjym/environments/popgym_count_recall.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/popjym/environments/popgym_count_recall.py b/popjym/environments/popgym_count_recall.py index 4943545..1e51593 100644 --- a/popjym/environments/popgym_count_recall.py +++ b/popjym/environments/popgym_count_recall.py @@ -24,12 +24,11 @@ class CountRecall(environment.Environment): def __init__(self, num_decks=1, num_types=2): super().__init__() self.decksize = 52 - # self.error_clamp = error_clamp # NOT IMPLEMENTED self.num_decks = num_decks self.num_types = num_types self.num_cards = self.decksize * self.num_decks self.max_num = self.num_cards // self.num_types - self.reward_scale = 1.0 / self.num_cards + self.reward_scale = 1.0 / (self.num_cards - 1) @property def default_params(self) -> EnvParams: @@ -39,15 +38,16 @@ def step_env( self, key: chex.PRNGKey, state: EnvState, action: int, params: EnvParams ) -> Tuple[chex.Array, EnvState, float, bool, dict]: - running_count = state.running_count.at[state.value_cards[state.timestep]].add(1) prev_count = state.running_count[state.query_cards[state.timestep]] + next_timestep = state.timestep + 1 + running_count = state.running_count.at[state.value_cards[next_timestep]].add(1) reward = jnp.where(action == prev_count, self.reward_scale, -self.reward_scale) new_state = EnvState( - state.timestep + 1, state.value_cards, state.query_cards, running_count + next_timestep, state.value_cards, state.query_cards, running_count ) obs = self.get_obs(new_state) - terminated = new_state.timestep == self.num_cards + terminated = new_state.timestep == self.num_cards - 1 return obs, new_state, reward, terminated, {} @@ -57,14 +57,13 @@ def reset_env(self, key: chex.PRNGKey, params: EnvParams) -> Tuple[chex.Array, E cards = jnp.arange(self.decksize * self.num_decks) % self.num_types value_cards = jax.random.permutation(key_value, cards) query_cards = jax.random.permutation(key_query, cards) - running_count = jnp.zeros((self.num_types,)) + running_count = jnp.zeros((self.num_types,)).at[value_cards[0]].add(1) state = EnvState( timestep=0, value_cards=value_cards, query_cards=query_cards, running_count=running_count, ) - # obs = state.cards[state.timestep] obs = self.get_obs(state) return obs, state @@ -77,7 +76,7 @@ def get_obs(self, state: EnvState) -> chex.Array: def action_space(self, params: Optional[EnvParams] = None) -> spaces.Discrete: """Action space of the environment.""" - return spaces.Discrete(2) + return spaces.Discrete(self.max_num + 1) def observation_space(self, params: EnvParams) -> spaces.Box: """Observation space of the environment.""" From e57620afe00b02958d658f4d826ad5e74a47be74 Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:34:11 +0200 Subject: [PATCH 07/14] fix: HigherLower off-by-one episode length causing OOB terminal read Episode ran one step too long, so the final step read a JAX-clamped out-of-bounds card. terminated now fires at new_state.timestep == num_cards - 1, matching popgym's len(deck) <= 1 termination and total step count. --- popjym/environments/popgym_higherlower.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/popjym/environments/popgym_higherlower.py b/popjym/environments/popgym_higherlower.py index dc7b5b7..eb95496 100644 --- a/popjym/environments/popgym_higherlower.py +++ b/popjym/environments/popgym_higherlower.py @@ -42,7 +42,7 @@ def step_env( reward = jnp.where(next_value == curr_value, 0, reward) new_state = EnvState(state.timestep + 1, state.cards) - terminated = new_state.timestep == num_cards + terminated = new_state.timestep == num_cards - 1 obs = self.get_obs(new_state) return obs, new_state, reward, terminated, {} From 1561203b2ce1e6c6a858dda7fde2581786708428 Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:34:14 +0200 Subject: [PATCH 08/14] fix: MineSweeper undersized obs, unreachable win, mine overwrite, timing - Observation size was sized to num_mines instead of min(num_mines + 1, 10), so on Easy the neighbor-count encoding silently dropped counts it couldn't represent. obs_size now matches popgym's Discrete(min(num_mines + 1, 10)) and is reused consistently for the grid, reset, and step logic. - Win condition checked a state that could never occur; corrected to jnp.all(new_grid != 0) (win once no CLEAR cells remain). - Revealing a mine overwrote the mine's own cell value instead of preserving it; jnp.where now keeps state.mine_grid on the mine's cell and only marks non-mine cells VIEWED. - Episode was one step too short; agent now gets the full max_episode_length steps. - Declare obs_requires_prev_action = True so make() auto-wraps the env with AliasPrevActionV2, matching popgym's contract for this env. --- popjym/environments/popgym_minesweeper.py | 26 +++++++++++++---------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/popjym/environments/popgym_minesweeper.py b/popjym/environments/popgym_minesweeper.py index 0b587e8..e37f8bc 100644 --- a/popjym/environments/popgym_minesweeper.py +++ b/popjym/environments/popgym_minesweeper.py @@ -21,6 +21,8 @@ class EnvParams: class MineSweeper(environment.Environment): + obs_requires_prev_action = True + def __init__(self, dims=(4, 4), num_mines=2): super().__init__() self.dims = dims @@ -29,6 +31,7 @@ def __init__(self, dims=(4, 4), num_mines=2): self.success_reward_scale = 1 / self.max_episode_length self.fail_reward_scale = -0.5 - self.success_reward_scale self.bad_action_reward_scale = -0.5 / (self.max_episode_length - 2) + self.obs_size = min(self.num_mines + 1, 10) @property def default_params(self) -> EnvParams: @@ -41,21 +44,24 @@ def step_env( mine = state.mine_grid[action] == 1 viewed = state.mine_grid[action] == 2 - new_grid = state.mine_grid.at[action].set(2) + new_grid = jnp.where( + mine, state.mine_grid, state.mine_grid.at[action].set(2) + ) reward = self.success_reward_scale reward = jnp.where(viewed, self.bad_action_reward_scale, reward) reward = jnp.where(mine, self.fail_reward_scale, reward) - terminated = state.timestep == self.max_episode_length + new_timestep = state.timestep + 1 + terminated = new_timestep > self.max_episode_length terminated = jnp.where(mine, True, terminated) - terminated = jnp.logical_or(terminated, jnp.all(new_grid == 2)) + terminated = jnp.logical_or(terminated, jnp.all(new_grid != 0)) - obs = jnp.zeros((self.num_mines,)) + obs = jnp.zeros((self.obs_size,)) obs = obs.at[state.neighbor_grid[action]].set(1) new_state = EnvState( - timestep=state.timestep + 1, + timestep=new_timestep, mine_grid=new_grid, neighbor_grid=state.neighbor_grid, ) @@ -64,7 +70,6 @@ def step_env( def reset_env(self, key: chex.PRNGKey, params: EnvParams) -> Tuple[chex.Array, EnvState]: """Performs resetting of environment.""" - # hidden_grid = jnp.zeros((params.dims[0] * params.dims[1],), dtype=jnp.int8) hidden_grid = jnp.zeros((self.dims[0] * self.dims[1],), dtype=jnp.int8) mines_flat = jax.random.choice( key, hidden_grid.shape[0], shape=(self.num_mines,), replace=False @@ -82,19 +87,18 @@ def reset_env(self, key: chex.PRNGKey, params: EnvParams) -> Tuple[chex.Array, E neighbor_grid=jnp.ravel(neighbor_grid), ) - return jnp.zeros((self.num_mines,)), state + return jnp.zeros((self.obs_size,)), state def action_space(self, params: Optional[EnvParams] = None) -> spaces.Discrete: """Action space of the environment.""" - # TODO: Multi-Discrete? return spaces.Discrete(np.prod(self.dims)) def observation_space(self, params: EnvParams) -> spaces.Box: """Observation space of the environment.""" return spaces.Box( - jnp.zeros((self.num_mines,)), - jnp.ones((self.num_mines,)), - (self.num_mines,), + jnp.zeros((self.obs_size,)), + jnp.ones((self.obs_size,)), + (self.obs_size,), dtype=jnp.float32, ) From e829196d5b4779715a46c0cd05d54a4dd2e2d68d Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:34:16 +0200 Subject: [PATCH 09/14] fix: MultiArmedBandit unlearnable without prev-action, empty info dict MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - The agent never saw its own previous action, so it had no way to associate a reward with the arm that produced it — the task was effectively unlearnable. Declare obs_requires_prev_action = True so make() auto-wraps the env with AliasPrevActionV2, matching popgym's contract for this env. - step() returned an empty info dict; now returns {"bandits": state.payouts} as popgym does. --- popjym/environments/popgym_multiarmedbandit.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/popjym/environments/popgym_multiarmedbandit.py b/popjym/environments/popgym_multiarmedbandit.py index 2417dcc..107ce45 100644 --- a/popjym/environments/popgym_multiarmedbandit.py +++ b/popjym/environments/popgym_multiarmedbandit.py @@ -19,6 +19,8 @@ class EnvParams: class MultiarmedBandit(environment.Environment): + obs_requires_prev_action = True + def __init__(self, num_bandits=10, episode_length=200): super().__init__() self.num_bandits = num_bandits @@ -41,7 +43,7 @@ def step_env( new_state = EnvState(timestep=state.timestep + 1, payouts=state.payouts) terminated = new_state.timestep >= self.episode_length - return obs, new_state, reward, terminated, {} + return obs, new_state, reward, terminated, {"bandits": state.payouts} def reset_env(self, key: chex.PRNGKey, params: EnvParams) -> Tuple[chex.Array, EnvState]: """Performs resetting of environment.""" From d8c559fe8dd1e3adfbc70de2eb85e9098a2ab47f Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:34:18 +0200 Subject: [PATCH 10/14] fix: RepeatFirst off-by-one episode length leaking the target card Episode ran one step too long. Combined with a defensive `% num_cards` wrap on the query index, the terminal step could wrap back around to the target card's own index, directly leaking the answer into the observation. terminated now fires at new_state.timestep == num_cards - 1 (matching popgym's step count), and the now-unreachable modulo wrap is removed since indices stay in-bounds. --- popjym/environments/popgym_repeat_first.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/popjym/environments/popgym_repeat_first.py b/popjym/environments/popgym_repeat_first.py index 301730c..b016b9a 100644 --- a/popjym/environments/popgym_repeat_first.py +++ b/popjym/environments/popgym_repeat_first.py @@ -37,7 +37,7 @@ def step_env( reward_scale = 1.0 / (num_cards - 1.0) reward = jnp.where(state.cards[0] == action, reward_scale, -reward_scale) new_state = EnvState(state.timestep + 1, state.cards) - terminated = new_state.timestep == num_cards + terminated = new_state.timestep == num_cards - 1 obs = self.get_obs(new_state) return obs, new_state, reward, terminated, {} @@ -57,7 +57,7 @@ def reset_env(self, key: chex.PRNGKey, params: EnvParams) -> Tuple[chex.Array, E def get_obs(self, state: EnvState) -> chex.Array: """Returns observation from the state.""" obs = jnp.zeros((self.num_suits,)) - obs = obs.at[state.cards[state.timestep % (self.decksize * self.num_decks)]].set(1) + obs = obs.at[state.cards[state.timestep]].set(1) return obs def action_space(self, params: Optional[EnvParams] = None) -> spaces.Discrete: From 9d66856486ada4537ea47849a94a4080a50ba62c Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:34:21 +0200 Subject: [PATCH 11/14] fix: RepeatPrevious recall lag off by one, off-by-one episode length - Effective recall lag was k instead of k - 1: the query index and the gate that enables querying were not shifted together. Query is now cards[state.timestep - self.k + 1] with gate state.timestep >= self.k - 1, so for k=4 the agent recalls cards[0] after observing cards[0..3] (lag = k - 1), matching popgym. - Episode was one step too long, which also caused a terminal out-of-bounds read; terminate at new_state.timestep == num_cards - 1. --- popjym/environments/popgym_repeat_previous.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/popjym/environments/popgym_repeat_previous.py b/popjym/environments/popgym_repeat_previous.py index 7ad5023..5d0ebbe 100644 --- a/popjym/environments/popgym_repeat_previous.py +++ b/popjym/environments/popgym_repeat_previous.py @@ -38,11 +38,11 @@ def step_env( reward = 0 reward = jnp.where( - state.cards[state.timestep - self.k] == action, reward_scale, -reward_scale + state.cards[state.timestep - self.k + 1] == action, reward_scale, -reward_scale ) - reward = jnp.where(state.timestep < self.k, 0, reward) + reward = jnp.where(state.timestep < self.k - 1, 0, reward) new_state = EnvState(state.timestep + 1, state.cards) - terminated = new_state.timestep == num_cards + terminated = new_state.timestep == num_cards - 1 obs = self.get_obs(new_state) return obs, new_state, reward, terminated, {} From 66b3702509ae9f7e0591bd0d135043ec99b147c8 Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:34:23 +0200 Subject: [PATCH 12/14] fix: CartPole spurious post-termination reward, unconditional noise clipping - The env emitted a -1.0 reward on the step after termination via a prev_terminal/reward_transform mechanism popgym's actual runtime behavior never exercises (its -1 branch is dead code under default gymnasium settings). reward is now unconditionally 1.0 / self.max_steps_in_episode per step, and the now-unused prev_terminal/reward_transform machinery is removed. - Stateless variants with noise_sigma == 0 were still being clipped after adding (zero) noise; clipping is now gated on self.noise_sigma > 0 so the noiseless variants return the raw [x, theta] observation as popgym does. --- popjym/environments/popgym_cartpole.py | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/popjym/environments/popgym_cartpole.py b/popjym/environments/popgym_cartpole.py index f1db2f1..3e197a6 100644 --- a/popjym/environments/popgym_cartpole.py +++ b/popjym/environments/popgym_cartpole.py @@ -52,7 +52,6 @@ def step_env( self, key: chex.PRNGKey, state: EnvState, action: int, params: EnvParams ) -> Tuple[chex.Array, EnvState, float, bool, dict]: """Performs step transitions in the environment.""" - prev_terminal = self.is_terminal(state, params) force = params.force_mag * action - params.force_mag * (1 - action) costheta = jnp.cos(state.theta) sintheta = jnp.sin(state.theta) @@ -71,9 +70,7 @@ def step_env( theta = state.theta + params.tau * state.theta_dot theta_dot = state.theta_dot + params.tau * thetaacc - # Important: Reward is based on termination is previous step transition - reward = 1.0 - prev_terminal - reward = self.reward_transform(params, reward) + reward = 1.0 / self.max_steps_in_episode # Update state dict and evaluate termination conditions state = EnvState(x, x_dot, theta, theta_dot, state.time + 1) @@ -87,9 +84,6 @@ def step_env( {"discount": self.discount(state, params)}, ) - def reward_transform(self, params: EnvParams, reward: float): - return jnp.where(jnp.isclose(reward, 0), -1.0, 1.0 / self.max_steps_in_episode) - def reset_env(self, key: chex.PRNGKey, params: EnvParams) -> Tuple[chex.Array, EnvState]: """Performs resetting of environment.""" init_state = jax.random.uniform(key, minval=-0.05, maxval=0.05, shape=(4,)) @@ -105,11 +99,14 @@ def reset_env(self, key: chex.PRNGKey, params: EnvParams) -> Tuple[chex.Array, E def get_obs(self, key: chex.PRNGKey, state: EnvState, params: EnvParams) -> chex.Array: """Applies observation function to state.""" - obs = ( - jnp.array([state.x, state.theta]) - + jax.random.normal(key, shape=(2,)) * self.noise_sigma - ) - obs = jnp.clip(obs, self.observation_space(params).low, self.observation_space(params).high) + obs = jnp.array([state.x, state.theta]) + if self.noise_sigma > 0: + obs = obs + jax.random.normal(key, shape=(2,)) * self.noise_sigma + obs = jnp.clip( + obs, + self.observation_space(params).low, + self.observation_space(params).high, + ) return obs def is_terminal(self, state: EnvState, params: EnvParams) -> bool: From aa39cbc3e29b2e64b2e4f30b94cd6c762df9c6c3 Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:34:27 +0200 Subject: [PATCH 13/14] fix: Pendulum noisy observation not clipped to bounds Noise added to the position-only observation could push it outside [-1, 1]; popgym clips via np.clip(obs + noise, obs_space.low, obs_space.high) after adding noise, which this port omitted. Apply jnp.clip(obs, -1.0, 1.0) after adding noise to match. --- popjym/environments/popgym_pendulum.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/popjym/environments/popgym_pendulum.py b/popjym/environments/popgym_pendulum.py index 1c88135..987253d 100644 --- a/popjym/environments/popgym_pendulum.py +++ b/popjym/environments/popgym_pendulum.py @@ -102,7 +102,7 @@ def reset_env(self, key: chex.PRNGKey, params: EnvParams) -> Tuple[chex.Array, E def get_obs(self, key: chex.PRNGKey, state: EnvState, params: EnvParams) -> chex.Array: """Return angle in polar coordinates and change.""" - return ( + obs = ( jnp.array( [ jnp.cos(state.theta), @@ -112,6 +112,7 @@ def get_obs(self, key: chex.PRNGKey, state: EnvState, params: EnvParams) -> chex ).squeeze() + jax.random.normal(key, shape=(2,)) * self.noise_sigma ) + return jnp.clip(obs, -1.0, 1.0) def is_terminal(self, state: EnvState, params: EnvParams) -> bool: """Check whether state is terminal.""" From a1d0118dab7797a6fc83649b8c7b29b6adc3e256 Mon Sep 17 00:00:00 2001 From: noahfarr Date: Tue, 7 Jul 2026 11:34:30 +0200 Subject: [PATCH 14/14] fix: meta_cartpole compat with newer JAX (jax.tree_map removed) jax.tree_map was removed in newer JAX releases in favor of jax.tree.map; update the one call site so this module imports on current JAX. --- popjym/environments/meta_cartpole.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/popjym/environments/meta_cartpole.py b/popjym/environments/meta_cartpole.py index fca6e72..965500e 100644 --- a/popjym/environments/meta_cartpole.py +++ b/popjym/environments/meta_cartpole.py @@ -60,7 +60,7 @@ def step_env( # env_obs_re, env_state_re = self.env.reset_env(key_reset, params.env_params) env_obs_re, env_state_re = state.init_obs, state.init_state - env_state = jax.tree_map( + env_state = jax.tree.map( lambda x, y: jax.lax.select(env_done, x, y), env_state_re, env_state_st ) env_obs = jax.lax.select(env_done, env_obs_re, env_obs_st)