Take the control mask into account when computing the loss in PolicyTask.
Take the control mask into account when computing the loss in PolicyTask.

PiperOrigin-RevId: 373219272

google/trax,google/trax