refactor: add missing train tags and only extract task data once

[m-muecke]

See @mllg PR changes for weights, this adds missing train tags, only calls task$data() once, and uses keys for reflection subsetting