Fix/env equivalence vs popgym

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)

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:

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

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:

popjym/environments/popgym_concentration.py

@@ -21,15 +21,17 @@ 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
         # 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.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)

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."""

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, {}

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,
         )
 

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."""

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."""