Add MarlGrid matrix observations (#424)

## Summary
- add semantic `matrix` and `full_matrix` observation formats for
MarlGrid
- keep existing pixel observations as the default via
`observation_format="pixels"`
- add direct MarlGrid tests for partial and full matrix observation
shapes and channels
- replace the `.dockerignore` symlink with explicit Docker ignore rules
- fix property docstrings that fail the current ruff lint rule

## Tests
- make bazel-test BAZEL_TEST_TARGETS=//envpool/marlgrid:marlgrid_test
- uvx ruff check envpool/python/envpool.py envpool/python/protocol.py
scripts/coverage_summary.py
- uvx ruff format --check envpool/python/envpool.py
envpool/python/protocol.py scripts/coverage_summary.py
envpool/marlgrid/marlgrid_test.py
- uvx clang-format --style=file -i envpool/marlgrid/marlgrid.h -n
--Werror
- git diff --check origin/main..HEAD

Note: the branch was rebuilt on top of current `origin/main` after
native MarlGrid prestige coloring landed upstream in #419.

---------

Co-authored-by: Jiayi Weng <jiayi@openai.com>

Author: tortoise10101

docs/env/marlgrid.rst

@@ -34,15 +34,40 @@ Options
 * ``prestige_beta (float)``: per-step prestige decay factor, default to
   ``0.95``;
 * ``prestige_scale (float)``: reward-history scale used when mapping prestige
-  to color, default to ``2.0``.
+  to color, default to ``2.0``;
+* ``observation_format (str)``: observation representation. ``"pixels"``
+  (default) returns the existing RGB partial view, ``"matrix"`` returns an
+  egocentric semantic partial view, and ``"full_matrix"`` returns the
+  world-oriented semantic grid.
 
 
 Observation Space
 -----------------
 
-MarlGrid returns one RGB partial-view image per player. The default registered
-tasks use ``view_tile_size=8`` and expose ``obs`` as a uint8 tensor with shape
-``(view_tile_size * view_size, view_tile_size * view_size, 3)`` per player.
+MarlGrid returns one uint8 observation per player. Its shape depends on
+``observation_format``:
+
+* ``pixels``: ``(view_tile_size * view_size, view_tile_size * view_size, 3)``.
+  This is the existing RGB partial view. The default registered tasks use
+  ``view_tile_size=8``;
+* ``matrix``: ``(view_size, view_size, 13)``. This is the agent's egocentric
+  partial view. Occluded cells and every cell for an inactive player are all
+  zero;
+* ``full_matrix``: ``(grid_size, grid_size, 13)``. This is the global grid in
+  world coordinates. Every cell for an inactive player is zero.
+
+For both matrix formats, channels are:
+
+* ``0``--``4``: one-hot empty, wall, goal, bonus, and lava base tiles;
+* ``5``: agent presence;
+* ``6``--``8``: agent red, green, and blue values. These contain the prestige
+  color when ``prestige_coloring=True``;
+* ``9``--``12``: one-hot agent direction: right, down, left, and up.
+
+One visible base-tile channel and every active one-hot channel use value
+``255``. An agent can occupy a base object, so base-tile and agent channels may
+both be set. Directions in ``matrix`` use the rotated egocentric coordinates;
+directions in ``full_matrix`` use world coordinates.
 
 Player metadata is returned under ``info["players"]``:
 

envpool/highway/highway_official_align_test.py

@@ -149,7 +149,7 @@ def _envpool_info(info: dict[str, Any]) -> dict[str, Any]:
 
 
 def _official_info(info: dict[str, Any]) -> dict[str, Any]:
-    exposed = {
+    exposed: dict[str, Any] = {
         key: info[key]
         for key in ("speed", "crashed", "is_success")
         if key in info

envpool/marlgrid/marlgrid.h

@@ -66,6 +66,10 @@ using Rgb = std::array<std::uint8_t, 3>;
 
 inline constexpr int kTilePixels = 32;
 inline constexpr int kTileSubdivs = 3;
+inline constexpr int kMatrixAgentChannel = 5;
+inline constexpr int kMatrixColorChannel = 6;
+inline constexpr int kMatrixDirectionChannel = 9;
+inline constexpr int kMatrixObsChannels = 13;
 inline constexpr double kPi = 3.14159265358979323846;
 inline constexpr Rgb kShadowColor = {35, 25, 30};
 inline constexpr Rgb kWorstColor = {74, 65, 42};
@@ -376,15 +380,32 @@ class MarlGridEnvFns {
                     "reset_on_mistake"_.Bind(false), "reward_decay"_.Bind(true),
                     "respawn"_.Bind(false), "ghost_mode"_.Bind(true),
                     "prestige_coloring"_.Bind(false),
-                    "prestige_beta"_.Bind(0.95f), "prestige_scale"_.Bind(2.0f));
+                    "prestige_beta"_.Bind(0.95f), "prestige_scale"_.Bind(2.0f),
+                    "observation_format"_.Bind(std::string("pixels")));
   }
 
   template <typename Config>
   static decltype(auto) StateSpec(const Config& conf) {
+    const std::string observation_format = conf["observation_format"_];
+    if (observation_format != "pixels" && observation_format != "matrix" &&
+        observation_format != "full_matrix") {
+      throw std::runtime_error(
+          "MarlGrid observation_format must be 'pixels', 'matrix', or "
+          "'full_matrix'");
+    }
     int obs_size = conf["view_size"_] * conf["view_tile_size"_];
+    int obs_channels = 3;
+    if (observation_format == "matrix") {
+      obs_size = conf["view_size"_];
+      obs_channels = detail::kMatrixObsChannels;
+    } else if (observation_format == "full_matrix") {
+      obs_size = conf["grid_size"_];
+      obs_channels = detail::kMatrixObsChannels;
+    }
     int bound = conf["grid_size"_];
     return MakeDict(
-        "obs"_.Bind(Spec<std::uint8_t>({-1, obs_size, obs_size, 3}, {0, 255})),
+        "obs"_.Bind(Spec<std::uint8_t>({-1, obs_size, obs_size, obs_channels},
+                                       {0, 255})),
         "info:players.id"_.Bind(Spec<int>({-1}, {0, conf["max_num_players"_]})),
         "info:players.done"_.Bind(Spec<bool>({-1})),
         "info:players.active"_.Bind(Spec<bool>({-1})),
@@ -423,6 +444,7 @@ class MarlGridEnv : public Env<MarlGridEnvSpec>, public RenderableEnv {
         prestige_coloring_(spec.config["prestige_coloring"_]),
         prestige_beta_(spec.config["prestige_beta"_]),
         prestige_scale_(spec.config["prestige_scale"_]),
+        observation_format_(spec.config["observation_format"_]),
         max_episode_steps_(spec.config["max_episode_steps"_]) {
     CHECK_GE(max_num_players_, n_agents_);
     CHECK_GE(grid_size_, 3);
@@ -546,6 +568,7 @@ class MarlGridEnv : public Env<MarlGridEnvSpec>, public RenderableEnv {
   bool prestige_coloring_{false};
   float prestige_beta_{0.95f};
   float prestige_scale_{2.0f};
+  std::string observation_format_{"pixels"};
   int max_episode_steps_{100};
   int step_count_{0};
   bool done_{true};
@@ -844,25 +867,29 @@ class MarlGridEnv : public Env<MarlGridEnvSpec>, public RenderableEnv {
     return detail::RenderAgentTile(agent.color, agent.dir, tile_size);
   }
 
+  [[nodiscard]] int TopAgentForCell(const detail::Cell* cell,
+                                    int top_agent_id) const {
+    if (cell == nullptr) {
+      return -1;
+    }
+    if (top_agent_id >= 0 &&
+        std::find(cell->agents.begin(), cell->agents.end(), top_agent_id) !=
+            cell->agents.end() &&
+        agents_[top_agent_id].active) {
+      return top_agent_id;
+    }
+    for (int agent_id : cell->agents) {
+      if (agents_[agent_id].active) {
+        return agent_id;
+      }
+    }
+    return -1;
+  }
+
   [[nodiscard]] std::vector<std::uint8_t> RenderTile(const detail::Cell* cell,
                                                      int top_agent_id,
                                                      int tile_size) const {
-    int chosen_agent = -1;
-    if (cell != nullptr) {
-      if (top_agent_id >= 0 &&
-          std::find(cell->agents.begin(), cell->agents.end(), top_agent_id) !=
-              cell->agents.end() &&
-          agents_[top_agent_id].active) {
-        chosen_agent = top_agent_id;
-      } else {
-        for (int agent_id : cell->agents) {
-          if (agents_[agent_id].active) {
-            chosen_agent = agent_id;
-            break;
-          }
-        }
-      }
-    }
+    int chosen_agent = TopAgentForCell(cell, top_agent_id);
     if (cell == nullptr || cell->type == detail::CellType::kEmpty) {
       if (chosen_agent >= 0) {
         auto tile = RenderAgentTile(chosen_agent, tile_size);
@@ -900,38 +927,41 @@ class MarlGridEnv : public Env<MarlGridEnvSpec>, public RenderableEnv {
     }
   }
 
-  void RenderAgentObs(int agent_id, std::uint8_t* output) const {
-    int obs_size = view_size_ * view_tile_size_;
-    std::vector<std::uint8_t> image(obs_size * obs_size * 3, 0);
-    for (int i = 0; i < obs_size * obs_size; ++i) {
-      image[i * 3 + 0] = detail::kShadowColor[0];
-      image[i * 3 + 1] = detail::kShadowColor[1];
-      image[i * 3 + 2] = detail::kShadowColor[2];
-    }
-    const detail::Agent& agent = agents_[agent_id];
-    if (!agent.active) {
-      std::memcpy(output, image.data(), image.size());
-      return;
-    }
+  [[nodiscard]] auto AgentViewCells(const detail::Agent& agent) const {
     int rot_k = (agent.dir + 1) % 4;
-    int orientation = (4 - rot_k) % 4;
     auto [top_x, top_y] = ViewTopLeft(agent);
-    std::vector<const detail::Cell*> view_cells(view_size_ * view_size_,
-                                                nullptr);
-    std::vector<bool> transparent(view_size_ * view_size_, true);
+    std::vector<const detail::Cell*> cells(view_size_ * view_size_, nullptr);
+    std::vector<bool> transparent(cells.size(), true);
     for (int y = 0; y < view_size_; ++y) {
       for (int x = 0; x < view_size_; ++x) {
         auto [sx, sy] = detail::RotateCoord(x, y, view_size_, rot_k);
         int wx = top_x + sx;
         int wy = top_y + sy;
         if (InBounds(wx, wy)) {
           const detail::Cell& cell = CellAt(wx, wy);
-          view_cells[detail::Offset(x, y, view_size_)] = &cell;
+          cells[detail::Offset(x, y, view_size_)] = &cell;
           transparent[detail::Offset(x, y, view_size_)] = cell.CanSeeBehind();
         }
       }
     }
-    std::vector<bool> visible = VisibilityMask(transparent);
+    return std::pair{std::move(cells), VisibilityMask(transparent)};
+  }
+
+  void RenderAgentObs(int agent_id, std::uint8_t* output) const {
+    int obs_size = view_size_ * view_tile_size_;
+    std::vector<std::uint8_t> image(obs_size * obs_size * 3, 0);
+    for (int i = 0; i < obs_size * obs_size; ++i) {
+      image[i * 3 + 0] = detail::kShadowColor[0];
+      image[i * 3 + 1] = detail::kShadowColor[1];
+      image[i * 3 + 2] = detail::kShadowColor[2];
+    }
+    const detail::Agent& agent = agents_[agent_id];
+    if (!agent.active) {
+      std::memcpy(output, image.data(), image.size());
+      return;
+    }
+    int orientation = 3 - agent.dir;
+    auto [view_cells, visible] = AgentViewCells(agent);
     for (int y = 0; y < view_size_; ++y) {
       for (int x = 0; x < view_size_; ++x) {
         if (!visible[detail::Offset(x, y, view_size_)]) {
@@ -953,6 +983,73 @@ class MarlGridEnv : public Env<MarlGridEnvSpec>, public RenderableEnv {
     std::memcpy(output, image.data(), image.size());
   }
 
+  static int MatrixObsOffset(int x, int y, int channel, int obs_size) {
+    return (y * obs_size + x) * detail::kMatrixObsChannels + channel;
+  }
+
+  void WriteMatrixBaseTile(const detail::Cell* cell, int x, int y, int obs_size,
+                           std::uint8_t* output) const {
+    int base_channel = cell == nullptr ? 0 : static_cast<int>(cell->type);
+    output[MatrixObsOffset(x, y, base_channel, obs_size)] = 255;
+  }
+
+  void WriteMatrixAgentTile(int chosen_agent, int x, int y, int obs_size,
+                            int orientation, std::uint8_t* output) const {
+    if (chosen_agent < 0) {
+      return;
+    }
+    detail::Rgb color =
+        prestige_coloring_
+            ? detail::PrestigeColor(agents_[chosen_agent].prestige,
+                                    prestige_scale_)
+            : detail::ColorValue(agents_[chosen_agent].color);
+    output[MatrixObsOffset(x, y, detail::kMatrixAgentChannel, obs_size)] = 255;
+    std::copy(
+        color.begin(), color.end(),
+        output + MatrixObsOffset(x, y, detail::kMatrixColorChannel, obs_size));
+    int direction = (agents_[chosen_agent].dir + orientation) % 4;
+    output[MatrixObsOffset(x, y, detail::kMatrixDirectionChannel + direction,
+                           obs_size)] = 255;
+  }
+
+  void WriteAgentMatrixObs(int agent_id, std::uint8_t* output) const {
+    int obs_values = view_size_ * view_size_ * detail::kMatrixObsChannels;
+    std::fill(output, output + obs_values, 0);
+    const detail::Agent& agent = agents_[agent_id];
+    if (!agent.active) {
+      return;
+    }
+    int orientation = 3 - agent.dir;
+    auto [view_cells, visible] = AgentViewCells(agent);
+    for (int y = 0; y < view_size_; ++y) {
+      for (int x = 0; x < view_size_; ++x) {
+        if (!visible[detail::Offset(x, y, view_size_)]) {
+          continue;
+        }
+        const detail::Cell* cell = view_cells[detail::Offset(x, y, view_size_)];
+        WriteMatrixBaseTile(cell, x, y, view_size_, output);
+        WriteMatrixAgentTile(TopAgentForCell(cell, agent_id), x, y, view_size_,
+                             orientation, output);
+      }
+    }
+  }
+
+  void WriteAgentFullMatrixObs(int agent_id, std::uint8_t* output) const {
+    int obs_values = grid_size_ * grid_size_ * detail::kMatrixObsChannels;
+    std::fill(output, output + obs_values, 0);
+    if (!agents_[agent_id].active) {
+      return;
+    }
+    for (int y = 0; y < grid_size_; ++y) {
+      for (int x = 0; x < grid_size_; ++x) {
+        const detail::Cell& cell = CellAt(x, y);
+        WriteMatrixBaseTile(&cell, x, y, grid_size_, output);
+        WriteMatrixAgentTile(TopAgentForCell(&cell, agent_id), x, y, grid_size_,
+                             0, output);
+      }
+    }
+  }
+
   void ResizeNearest(const std::uint8_t* src, int src_width, int src_height,
                      std::uint8_t* dst, int dst_width, int dst_height) const {
     for (int y = 0; y < dst_height; ++y) {
@@ -975,7 +1072,14 @@ class MarlGridEnv : public Env<MarlGridEnvSpec>, public RenderableEnv {
       state["info:players.pos"_](i, 1) = agents_[i].y;
       state["info:players.dir"_][i] = agents_[i].dir;
       state["reward"_][i] = last_rewards_[i];
-      RenderAgentObs(i, static_cast<std::uint8_t*>(state["obs"_][i].Data()));
+      auto* obs = static_cast<std::uint8_t*>(state["obs"_][i].Data());
+      if (observation_format_ == "matrix") {
+        WriteAgentMatrixObs(i, obs);
+      } else if (observation_format_ == "full_matrix") {
+        WriteAgentFullMatrixObs(i, obs);
+      } else {
+        RenderAgentObs(i, obs);
+      }
     }
   }
 };