Fix/env equivalence vs popgym#1
Open
noahfarr wants to merge 14 commits into
Open
Conversation
…ency 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.
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.
…ength - 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.
- 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.
- 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.
… 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.
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.
- 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.
- 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.
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.
- 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.
…lipping - 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.
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.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
I went through each of the 11 environments and compared them against the original gymnasium implementations in proroklab/popgym, since I'm relying on this package for a downstream project and wanted to be sure the JAX port actually matches the reference task. It doesn't in a few places and some of these change the actual difficulty or even solvability of the task.. This PR fixes what I found, one environment per commit so it's easy to review or drop anything you disagree with.
Autoencode: the WATCH/PLAY observation was inverted: the agent was shown the card during PLAY and zeros during WATCH, which is backwards. There was also an out-of-bounds read once the timestep passed num_cards, and the episode ran one step too long.
Battleship: ship placement sampled directions non-uniformly, biasing ships toward the top-left of the board instead of being uniform across all 4 directions.
CountRecall: this one was badly broken: action_space was hardcoded to Discrete(2) no matter the difficulty, so Medium and Hard were literally unsolvable since the agent couldn't output the right range of actions. On top of that: the previous-count value was read at the wrong point in the step, the running count wasn't seeded with the first dealt card on reset, the reward scale used the wrong denominator, and the episode ran one step too long.
HigherLower: off-by-one episode length, causing an out-of-bounds read on the final step.
MineSweeper: the observation array was sized to num_mines instead of min(num_mines + 1, 10), so on Easy the neighbor-count encoding silently dropped counts it had no room to represent. The win condition also checked a state that could never actually occur, and revealing a mine overwrote the mine's own cell value instead of preserving it.
MultiArmedBandit: the agent never got to see its own previous action, so it had no way to tell which arm produced which reward. This made the task effectively unlearnable unless the user themselves add an action embedding to the observation. In popgym they apply a wrapper automatically to the environments that need the action in their observation.
RepeatFirst: same off-by-one episode length pattern as the other card games.
RepeatPrevious: the effective recall lag was k instead of k - 1, and the episode again ran one step too long.
CartPole: a -1.0 reward was emitted on the step right after termination, which the reference never actually produces under default gymnasium settings.
Pendulum: noisy observations weren't clipped back into [-1, 1] after adding noise, unlike the reference.
Four environments (Battleship, Concentration, MineSweeper, MultiArmedBandit) rely on the reference's obs_requires_prev_action contract, which wasn't being honored at all. I added an opt-in mechanism in registration.py: envs that declare obs_requires_prev_action = True now get auto-wrapped with AliasPrevActionV2 in make(), with a printed notice so it's obvious when it kicks in.