diff --git a/lrax/common/activations.py b/lrax/common/activations.py index 5c6e7e4..7c6a620 100644 --- a/lrax/common/activations.py +++ b/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 ---------- diff --git a/lrax/common/networks.py b/lrax/common/networks.py index 5096f2d..8ba891e 100644 --- a/lrax/common/networks.py +++ b/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 ---------- diff --git a/lrax/ppo/ppo.py b/lrax/ppo/ppo.py index b7918d6..892332d 100644 --- a/lrax/ppo/ppo.py +++ b/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 ----- diff --git a/lrax/sac/sac.py b/lrax/sac/sac.py index d6cd287..684fb00 100644 --- a/lrax/sac/sac.py +++ b/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 diff --git a/lrax/shac/shac.py b/lrax/shac/shac.py index 3aa2343..a201d84 100644 --- a/lrax/shac/shac.py +++ b/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