Refactor rule decoding to remove lambda functions

[adzcai]

I noticed the comment in rules.py about evaluating all rules. I think this is unavoidable under vmap, but in the single-threaded case, by passing the operands in the way below, Jax should evaluate only the selected one:

# old
grid, agent = jax.lax.switch(
    encoding[0],
    (
        lambda: EmptyRule.decode(encoding)(grid, agent, action, position),
        ...
    ),
)

# new
grid, agent = jax.lax.switch(
    encoding[0],
    (
        EmptyRule.decode(encoding),
        ...
    ),
    grid,
    agent,
    action,
    position,
)

Let me know the best way to benchmark this!

EDIT: Sorry, upon further thought, I was hasty, think these might be exactly the same. Feel free to close this.

[Howuhh]

hi @adzcai!

You can benchmark times with the help of scripts/benchmark_xland.py and scripts/benchmark_xland_all.py scripts.