Traffic junction

[lucaschoii]

Description

This PR adds the Traffic Junction environment to JaxMARL, a fully JAX-compatible port of the classic multi-agent traffic control task.

This implementation features a scalable grid size, dynamic spawning logic, and a custom Matplotlib visualizer for evaluating agent behaviors.

Key Features

• Scalable Environment: The TrafficJunction class supports variable grid_size (default: 14), automatically adjusting lane positions (mid_low, mid_high) and spawn points.
• Right-Hand Drive (RHD) Logic: Implements correct pivot logic for RHD turns, preventing head-on collisions and ensuring realistic traffic flow.
• Observation Space: Agents receive a local view (view_size x view_size) centered on themselves, padded with -1 for off-grid areas.
• Visualization: Includes tj_visualizer.py with:
  • Dynamic marker scaling based on grid_size.
  • Directional triangle markers for moving agents.
  • Octagon markers for stopped agents (Stop Signs).
  • Collision and active agent counters.

Files Added

• jaxmarl/environments/traffic_junction/traffic_junction.py: Core environment logic.
• jaxmarl/environments/traffic_junction/tj_visualizer.py: Visualization tool.
• jaxmarl/environments/traffic_junction/test_tj.py: Sanity check / test script.

Verification

• Verified zero-collision spawning with off-grid buffers.
• Verified correct turning logic for all 4 directions (Right, Down, Left, Up).
• Verified visualizer scaling for different grid sizes.

[Copilot]

Pull request overview

Adds a new Traffic Junction multi-agent environment to JaxMARL (JAX-native), along with a Matplotlib-based visualizer and a small runnable sanity-check script.

Changes:

• Introduces TrafficJunction environment with grid-size-dependent lane/spawn logic, turning rules, and local observations.
• Adds TrafficVisualizer to animate recorded state sequences.
• Adds test_tj.py as a manual simulation/visualization driver.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 9 comments.

File	Description
jaxmarl/environments/traffic_junction/traffic_junction.py	Core environment implementation: state, stepping (movement/spawn), rewards, and observations.
jaxmarl/environments/traffic_junction/tj_visualizer.py	Matplotlib animation utility for inspecting state sequences over time.
jaxmarl/environments/traffic_junction/test_tj.py	Manual script to run random actions, track collisions, and launch the visualizer.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 6 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.