- [x] train state - [x] class - [x] pydantic - [x] create train state - [x] n_jitted steps - [x] get action - [x] transition - [ ] hyperparameter - [ ] performance report