The TODO in the forward of PGPolicy about renaming logits should be addressed. It will likely affect the naming in ModelOutputBatchProtocol and in other places, as well as docstrings. The exact usage of the model outputs should be clarified everywhere (they are passed to a torch Distribution which is then used for sampling or for computing log_prob).

Originally posted by @Trinkle23897 in #912 (comment)