Standardize citation formatting in docstrings

lrax/common/activations.py

@@ -25,11 +25,12 @@
 def lipswish(x: Array) -> Array:
     """LipSwish activation function.
 
-    LipSwish has been shown to perform well when applied to neural SDEs. Refer to the
-    following papers for additional information:
+    LipSwish has been shown to perform well when applied to neural SDEs. See:
 
-    - "Residual Flows for Invertible Generative Modeling" (https://arxiv.org/abs/1906.02735)
-    - "Efficient and Accurate Gradients for Neural SDEs" (https://arxiv.org/pdf/2105.13493)
+    - "Residual Flows for Invertible Generative Modeling", Ricky T. Q. Chen, Jens
+      Behrmann, David Duvenaud, and Jörn-Henrik Jacobsen.
+    - "Efficient and Accurate Gradients for Neural SDEs", Patrick Kidger, James Foster,
+      Xuechen Li, and Terry Lyons.
 
     Parameters
     ----------

lrax/common/networks.py

@@ -116,10 +116,7 @@ def __call__(self, x: Array) -> Array:
 class TamedMLP(eqx.Module):
     """Implements a tamed MLP, which helps prevent model blow-up during training.
 
-    For additional information, see,
-
-    "On neural differential equations"
-    (https://ora.ox.ac.uk/objects/uuid:af32d844-df84-4fdc-824d-44bebc3d7aa9)
+    "On Neural Differential Equations", Patrick Kidger.
     """
 
     mlp: eqx.nn.MLP
@@ -136,7 +133,7 @@ def __init__(
         *,
         key: PRNGKeyArray,
     ):
-        """Initialize an SPD network.
+        """Initialize a tamed MLP.
 
         Parameters
         ----------

lrax/ppo/ppo.py

@@ -70,7 +70,7 @@ def _compute_gae(
     """Calculate the advantages and returns using generalized advantage estimation.
 
     "High-Dimensional Continuous Control Using Generalized Advantage Estimation",
-    John Schulman, Philipp Moritz, Sergey Levine, Michael I. Jordan and Pieter Abbeel
+    John Schulman, Philipp Moritz, Sergey Levine, Michael I. Jordan, and Pieter Abbeel.
 
     Parameters
     ----------
@@ -140,7 +140,7 @@ class PPO(AbstractAlgorithm):
     """Proximal Policy Optimization.
 
     "Proximal Policy Optimization Algorithms", John Schulman, Filip Wolski, Prafulla
-    Dhariwal, Alec Radford, Oleg Klimov.
+    Dhariwal, Alec Radford, and Oleg Klimov.
 
     Notes
     -----

lrax/sac/sac.py

@@ -55,7 +55,7 @@ class SAC(AbstractAlgorithm):
     """Soft Actor-Critic.
 
     "Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a
-    Stochastic Actor", Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, Sergey Levine.
+    Stochastic Actor", Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine.
     """
 
     buffer_size: int = 1_000_000

lrax/shac/shac.py

@@ -141,7 +141,7 @@ class SHAC(AbstractAlgorithm):
 
     "Accelerated Policy Learning with Parallel Differentiable Simulation", Jie Xu,
     Viktor Makoviychuk, Yashraj Narang, Fabio Ramos, Wojciech Matusik, Animesh Garg,
-    Miles Macklin.
+    and Miles Macklin.
     """
 
     num_steps: int = 32