Jax - vmap over batch of dataclasses

In JAX, I am looking to vmap a function over a fixed length list of dataclasses, for example:

import jax, chex
from flax import struct

@struct.dataclass
class EnvParams:
    max_steps: int = 500
    random_respawn: bool = False

def foo(params: EnvParams):
    ...

param_list = jnp.Array([EnvParams(max_steps=500), EnvParams(max_steps=600)])
jax.vmap(foo)(param_list)

The example above fails since is not possible to create a jnp.Array of custom objects, and JAX doesn't allow vmapping over Python Lists. The only remaining option I see is to transform the dataclass to represent a batch of parameters, as so:

@struct.dataclass
class EnvParamBatch:
    max_steps: jnp.Array = jnp.array([500, 600])
    random_respawn: jnp.Array = jnp.array([False, True])

def bar(params):
    ...

jax.vmap(bar)(EnvParamBatch())

It would be preferable to use a container of structs (with each representing a single parameter set), so I'm wondering if there are any alternative approaches to this?

N.B. I am aware of this answer, however it's not precisely the same question and there may now be better solutions.