Tricky NonConcreteBooleanIndexError

[harryjulian]

I'm writing a small RL environment and getting a tricky NonConcreteBooleanIndexError when I'm calculating a boolean from an array. I wrote the initial version using python control flow but I've migrated the majority of my step function to jax.lax.cond except from this particular conditional below

DECK = list(range(0, 13)) # all possible card values
DTYPE = jnp.int8

...

# Compute mask of valid actions
card_vals = jnp.array(DECK, dtype=DTYPE)
available_mask = jnp.where(agent.hand.held, True, False)
legal_mask = jnp.where(
    (card_vals >= state.to_beat) | (card_vals == TWO) | (card_vals == THREE) | (card_vals == TEN), 
    True,
    False
)
valid_mask = available_mask & legal_mask

# Peform action selection
if jnp.any(agent.hand.held):
    action, state_update = _select_from_hand(key, agent, state, state_update, valid_mask)
else:
    action, state_update = _select_from_undercarriage(key, agent, state, state_update, valid_mask)

agent.hand.held is of shape (13,) of counts of a particular value of card

I've attempted to move to this,

action, state_update = jax.lax.cond(
    jnp.any(agent.hand.held),
    _select_from_hand,
    _select_from_undercarriage,
    key, agent, state, state_update, valid_mask
)

But keep receiving this error which corresponds to the valid_mask argument:

Cell In[28], line 202, in ShitheadEnvironment.step(self, key, state)
    199 else:
    200     action, state_update = _select_from_undercarriage(key, agent, state, state_update, valid_mask)
--> 202 action, state_update = jax.lax.cond(
    203     jnp.any(agent.hand.held),
    204     _select_from_hand,
    205     _select_from_undercarriage,
    206     key, agent, state, state_update, valid_mask
    207 )
    210 # Determine if we need to clear the table after a ten
    211 state_update = jax.lax.cond(
    212     action == TEN,
    213     _clear_table,
    214     lambda state_update: state_update,
    215     state_update
    216 )
...
   6978   out.append(bool(i))

NonConcreteBooleanIndexError: Array boolean indices must be concrete; got ShapedArray(bool[13])

I understand that shapes are static and need to be known at compile-time in JAX., but the valid_mask is a mask over agent.hand.held and will always by the same shape as it. I'm stuck!

I've tried a few extra things like this too, to no avail

action, state_update = jax.lax.cond(
    jnp.where(jnp.sum(agent.hand.held) > 0, True, False),
    _select_from_hand,
    _select_from_undercarriage,
    key, agent, state, state_update, valid_mask
)

[jakevdp]

It sounds like _select_from_undercarriage and/or _select_from_hand require valid_mask to be static, but valid_mask is not static because it is the result of a JAX operation. You'll probably need to update your functions to allow for dynamic values of valid_mask.

From the snippet of the traceback you shared, it may be as simple as changing out.append(bool(i)) to out.append(i) (though that's just a guess based on very incomplete information about your code).

[harryjulian]

Thanks for the speedy reply! The bottom of the traceback is from another library unfortunately.

I thought the shape of valid_mask should be static? It's computed based on the shape of card_vals and agent.hand.held which have the same shape (13,) as determined by some constants a little earlier on in the programme.

(I understand this is hard to reason about without seeing the whole file)

[jakevdp]

The shape is static, but the contents are not. The non-static contents are why you needed to switch from if valid_mask.any() to lax.cond(valid_mask.any(), ....

This was the right fix for that particular issue, but it revealed a deeper issue within the definitions of _select_from_hand and/or _select_from_undercarriage, and you haven't included enough information in your question for us to help you debug that.

[harryjulian]

That makes sense. Here's the additional funcs if you can see anything immediately wrong with them. There's another layer of nesting which could be at fault, too.

def _select_from_hand(key, agent, state, state_update, valid_mask):

    # Choose action
    x = jnp.concatenate([agent.hand.held, agent.hand.under, state.on_the_table])
    action_dist, key = agent.network(x, key)
    action = jnp.where(jnp.any(valid_mask), jnp.argmax(action_dist * valid_mask), -1)

    # Get state update based on valid action, when in hand we don't play if there are invalid actions
    state_update = jax.lax.cond(
        action != -1,
        _handle_succesful_hand_turn,
        _handle_table_pickup, 
        agent, action, state, state_update
    )

    return action, state_update
  
def _select_from_undercarriage(key, agent, state, state_update, valid_mask):

    # Choose random undercarriage card
    key, key2 = jax.random.split(key)
    indices = jnp.arange(0, N_CARD_TYPES, dtype=DTYPE)[agent.hand.under > 0]
    action = jax.random.choice(key2, indices, shape=(1,)).item()

    # Play it regardless of action validity
    state_update["under"] = agent.hand.under.at[action].add(-1)
    state_update["on_the_table"] = state.on_the_table.at[action].add(1)

    # Assess action validity and if invalid, pickup else continue
    state_update = jax.lax.cond(
        action >= state.to_beat,
        lambda *, state_update: state_update, # null condition
        _handle_table_pickup,
        agent, state, state_update
    )

    return action, state_update

Thinking about it I imagine the indices line will also cause a problem for the JIT, but shouldn't be the source of this first problem.

[jakevdp]

The problem is in this line:

indices = jnp.arange(0, N_CARD_TYPES, dtype=DTYPE)[agent.hand.under > 0]

since agent.hand.under is traced, there's no way to know statically how many of its values are larger than zero, and so indices would have to have a dynamic size, and dynamically-sized arrays are not supported in JAX transforms like jit (see JAX sharp bits: Dynamic Shapes).

You could probably work around this by using jax.random.randint rather than jax.random.choice, and avoiding constructing the explicit indices array.

[harryjulian]

Awesome, thanks Jake.