jax.jit しやすいシミュレータの実装方針

[sotetsuk]

• [必須] 可変長のforループを使わない
• [必須] listを使わない（appendなども使わない）
• ifのネストやforをなるべくさけ、numpyの操作をなるべく使う
• 長い関数や深いネストをさけ、10行以内くらいの小さいpure functionになるべく分割する
  • pure functionとは副作用が起こらず、引数が同じなら同じ結果が返ってくるもののこと
• stateが一つの行列で表現できない場合は dataclass で実装する
  • dataclassの各要素はすべて np.ndarray とし、intの場合も a: np.ndarray = np.zeros(1) といった具合にする
  • observationを生成する関数も用意する
• actionがintや一つのベクトルで表現できない場合は、 dataclass で実装する Actionは一つのintで表現できる必要がある（legal_actionsが可変長リストになってしまうため）
  • 0~N個の連続した数字に各アクションを割り当てる必要がある。
• legal_actionsも可変長リストはダメ。
  • ダメな例: legal_actions: List[Action] みたいなのはダメ。