v3.0.0

Kinetix v3.0.0

This release adds a number of new features, the primary ones being:

1. Data loading utilities and offline BC training for the offline dataset
2. Shard mapped training code for ppo.py and sfl.py, allowing training to scale to multiple GPUs and multiple nodes.

---

New features

• Offline BC training (experiments/offline_bc.py): shard-mapped training with warmup LR schedule, validation, checkpointing, and GIF rendering
• Data loading (kinetix/data/): Grain-compatible Zarr data sources (ZarrBatchDataSource, ZarrTrajDataSource, MultiFileBatchDataSource); see examples/example_data_loading.py
• RMS normalisation (kinetix/util/learning.py): RunningMeanStandard, rms_init, rms_normalise, parallel_rms_update; TrainStateWithRMSNorm carries and checkpoints the norm state; opt-in via rms_norm: true in config
• Model updates: TemperatureCategorical, actor_depth/actor_width params, new size presets configs/model/tf-{s,m,l,paper}.yaml

---

Breaking changes

Renamed fc_layer_width and fc_layer_depth in the configs to actor_depth/critic_depth and actor_width/critic_width

v2.0.0

What's Changed

• Added error message when num minibatches is larger than num train envs by @Michael-Beukman in #24
• Refactor kinetix.environment to have a separate spaces module to avoid circular imports in the future. by @Michael-Beukman in #21
  • This is a breaking change since many of the imports changed.

Full Changelog: v1.0.5...v2.0.0

v1.0.5

bump jax & flax version

v1.0.2

What's Changed

• Added an optional flag to the make_env function that controls if we d… by @Michael-Beukman in #15
• Consistently use jax.tree.map by @Michael-Beukman in #16

Full Changelog: v1.0.0...v1.0.2

v1.0.0

What's Changed

• Kinetix v1.0.0 by @Michael-Beukman in #10

Full Changelog: v0.1.0...v1.0.0

v0.1.0

v0.1.0 Release of Kinetix, reproducing the results of the paper.