Skip to content

# SMAX

## Description

SMAX is a purely JAX SMAC-like environment. It, like SMAC, focuses on decentralised unit micromanagement across a range of scenarios. Each scenario features fixed teams.

## Scenarios

Name | Ally Units | Enemy Units
--- | --- | ---
2s3z | 2 stalkers & 3 zealots | 2 stalkers & 3 zealots
3s5z | 3 stalkers & 5 zealots | 3 stalkers & 5 zealots
5m_vs_6m | 5 marines | 6 marines
10m_vs_11m | 10 marines | 11 marines
27m_vs_30m | 27 marines | 30 marines
3s5z_vs_3s6z | 3 stalkers & 5 zealots | 3 stalkers & 6 zealots
3s_vs_5z | 3 stalkers | 5 zealots
6h_vs_8z | 6 hydralisks | 8 zealots
smacv2_5_units | 5 randomly chosen | 5 randomly chosen
smacv2_10_units | 10 randomly chosen | 10 randomly chosen
smacv2_20_units | 20 randomly chosen | 20 randomly chosen

## Visualisation

You can see the example `smax_introduction.py` in the tutorials folder for an introduction to SMAX, including example visualisation. SMAX environments tick at 8 times faster than each step of the agent. This means that when visualising, we have to expand the state sequence to encompass all ticks. This is why the `state_seq` for SMAX consists of a sequence of `(key, state, actions)` -- we must have not only the state and actions, but also the exact key passed to the step function to interpolate between the different states correctly. This process means visualisation can be time consuming if done for a large number of steps.

    from jaxmarl import make
    from jaxmarl.environments.smax import map_name_to_scenario
    from jaxmarl.viz.visualizer import SMAXVisualizer

    scenario = map_name_to_scenario("3m")
    env = make(
        "HeuristicEnemySMAX",
        enemy_shoots=True,
        scenario=scenario,
        num_agents_per_team=3,
        use_self_play_reward=False,
        walls_cause_death=True,
        see_enemy_actions=False,
    )

    # state_seq is a list of (key_s, state, actions) tuples
    # where key_s is the RNG key passed into the step function,
    # state is the jax env state and actions is the actions passed
    # into the step function.
    viz = SMAXVisualizer(env, state_seq)

    viz.animate(view=False, save_fname="output.gif")