google-deepmind / google-deepmind/acme
Improving snapshotting in JAX distributed experiments
- Dominant language
- Python
- Stars
- 4.1k
- Forks
- 553
- PR merge metrics
- No merged PRs in 30d
Description
Hi,
I have started adopting `acme.jax.experiments.make_distributed_experiment` for training distributed agents with LaunchPad. Right now, it seems that there is no way for me to configure snapshotting, which is a useful utility, and I was wondering if there is interest in improving this part from the Acme developers. I have a few proposals and I am happy to make a PR which includes some improvements. Here they are
## Support snapshotting configuration.
Right now it's not possible to configure anything about the snapshotting. It would be great if we include options to specify
1) where to store snapshots. It currently stores the snapshots in the workdir, without any subdirectory structure so you get workdir/20220601-00000/. At the bare minimum by default, we should put them in a subdirectory so that we have workdir/snapshots/20220601-00000/
2) how frequent to snapshot. For long training jobs (I work with pixels where training on humanoid control may take a day), the current default frequency is 5 minutes, and I end up with hundreds of snapshots that I don't really need.
To ensure backward compatibility so that users can opt-in to the new improvement. We can fall back to the same configuration as now if the user does not specify the configuration for snapshotting.
## Improve the snapshotting format
Right now, JAX ModelToSnapshot embeds parameters as fixed constants in the tf.SavedModel. This is less general especially if users want to finetune the models afterward. Maybe we can save the parameters as variables in the SavedModel? I saw that the TF snapshotter is doing that and we can adopt a similar fashion, similar to the way that was done in https://github.com/google/jax/tree/main/jax/experimental/jax2tf
## Add utilities for inference with snapshots.
I am not sure how users use the snapshots internally, perhaps with some internal libraries which allows easy loading of saved policies for evaluation/rollouts, but it would be great if there is at least some bare minimum way to do this in the OSS version. To me, the snapshots are really an excellent way to store models for analysis/evaluation/data generation afterwards. For example, I use the saved model for generating expert trajectories which I then use for imitation learning. It would be great if there are some helper functions that makes this work easily, or at least some examples of how users of Acme can take advantage of the saved snapshots.
Best,
Yicheng
Contributor guide
Assessment
This issue has not been assessed yet.