google-deepmind / google-deepmind/acme
Is the last layer in distributional critic optimized (jax implementation)?
Open
- Dominant language
- Python
- Stars
- 4.1k
- Forks
- 553
- PR merge metrics
- No merged PRs in 30d
Description
Hi! I was wondering if the last layer in the distributional critic used in dmpo, for example, is optimized in the jax implementation, since it is instantiated in the __call__ function and not in the __init__?
https://github.com/deepmind/acme/blob/98c4204b9be7f327035bf1dbef26aa820cc4c1ec/acme/jax/networks/distributional.py#L373-L414
Contributor guide
Assessment
This issue has not been assessed yet.