Suggestion on probs parameter of categorical dist in numpyro

I have a question posted many days, the model is quite straight forward, would you please spend some time taking a look at it?
Sorry for asking too much questions~ I’m quite new to pyro/numpyro but my work is somewhat difficult for my current ability.