Work on documentation (17). Fix some typing errors

- Finish PyTrees example.
- Fix a few more typing errors.
- Bind `StateT` to `PyTree[Array]` rather than `Module, Array`
  or just `Module`.
- Edit `docs/index.md`. Move issue categories to a GitHub
  conversation
- Remove `RNNCellProto`.

docs/_overrides/.icons/material/partials/comments.html

@@ -3,7 +3,7 @@ <h2 id="__comments">{{ lang.t("meta.comments") }}</h2>
   <script src="https://giscus.app/client.js"
     data-repo="mlprt/feedbax"
     data-repo-id="R_kgDOKLhZBg"
-    data-category="Announcements"
+    data-category="Documentation"
     data-category-id="DIC_kwDOKLhZBs4Cdi3X"
     data-mapping="pathname"
     data-strict="0"
@@ -28,7 +28,7 @@ <h2 id="__comments">{{ lang.t("meta.comments") }}</h2>
         : "light"
 
       // Instruct Giscus to set theme
-      giscus.setAttribute("data-theme", theme) 
+      giscus.setAttribute("data-theme", theme)
 
 
     }

docs/api/networks.md

@@ -6,7 +6,4 @@
 
 ::: feedbax.nn.SimpleStagedNetwork
 
-::: feedbax.nn.RNNCellProto
-    options:
-        members: ['__init__', '__call__'] 
 

docs/index.md

@@ -1,25 +1,26 @@
 # Getting started
 
-Feedbax is a [JAX](https://jax.readthedocs.io/en/latest/beginner_guide.html#beginner-guide) library for optimal (feedback) control.
+Feedbax is a [JAX](https://jax.readthedocs.io/en/latest/beginner_guide.html#beginner-guide) library for optimal feedback control with neural networks.
 
 Feedbax makes it easy to:
 
-- [train](/feedbax/examples/0_train_simple) neural networks to control the movement of simulated limbs (biomechanical models);
+- [train](/feedbax/examples/0_train_simple) a neural network to control a simulated limb (biomechanical model) to perform movement tasks;
 - [intervene](/feedbax/examples/3_intervening) on existing models and tasks—for example, to:
-    - add force fields to trials at random;
+    - add force fields that disturb a limb;
     - alter the activity of a single unit in a neural network;
     - perturb the sensory feedback received by a network;
     - add any kind of noise to any part of a model's state;
-- [train replicates](/feedbax/examples/4_vmap) of a model in parallel;
-- specify which parts of the model are [trainable](/feedbax/examples/1_train/#selecting-part-of-the-model-to-train), or available to a controller as feedback;
+- schedule an intervention to occur on only a subset of task trials or time steps;
+- specify which parts of the model are [trainable](/feedbax/examples/1_train/#selecting-part-of-the-model-to-train), and which parts are available as feedback to a neural network;
+- train [multiple replicates](/feedbax/examples/4_vmap) of a model at once;
 - swap out components of models, and write new components.
 <!-- - track the progress of a training run in Tensorboard. -->
 
 Feedbax was designed for feedback control of biomechanical models, to perform movement tasks in continuous spaces. However, it can be used more generally as a framework for optimal control.
 
 ## Feedbax is a JAX library
 
-Feedbax uses JAX and [Equinox](https://docs.kidger.site/equinox/), because their features are very convenient for model design and scientific analysis.
+Feedbax uses JAX and [Equinox](https://docs.kidger.site/equinox/).
 
 [Never used JAX before](/feedbax/examples/pytrees/)?
 <!--
@@ -29,28 +30,22 @@ For a library that's similar to Feedbax but written in PyTorch, please check out
 
 ## Installation
 
-Pip TODO.
+`pip install feedbax`
 
-`python -m pip install`
-
-### Installing from source
+Currently requires Python>=3.11.
 
 ## Development
 
-I've developed Feedbax over the last few months, as I've learned JAX. My short-term objective has been to serve my own use case—graduate research in the neuroscience of motor control—but I have also tried to make design choices in pursuit of reusability and generality.
+I've developed Feedbax over the last few months, while learning JAX. My short-term objective has been to serve my own use case—graduate research in the neuroscience of motor control—but I've also tried to design something reusable and general.
 
-By making the library open source now, I hope to receive some feedback about those decisions. To make that easier I've created GitHub [issues](https://github.com/mlprt/feedbax/issues) documenting my choices and uncertainties. Those issues largely fall into a few categories:
+I've added GitHub [issues](https://github.com/mlprt/feedbax/issues) to document some of my choices and uncertainties, and to open them for discussion. For an overview of major issues in different categories, check out [this GitHub conversation](https://github.com/mlprt/feedbax/discussions/27).
 
-1. Structure: Some of the abstractions I've chosen are probably clumsy. It would be good to find out, at this point—maybe we can make some changes for the better!
-   In approximate order of significance: #19, #24, #12, #20, #1, #5, #21.
-2. Features: There are many small additions that could be made, especially to the pre-built models and tasks. There are also a few major improvements which I am anticipating in the near future, such as *online learning* (#21). #10,
-3. Typing: I am currently working through many pyright errors, trying to make Feedbax compliant. I'm not very experienced with typing in Python, so I may have tried to do things that are heavy-handed or over-clever. See issues: (#7, #8, #9, #11)
-4. Performance.
+There are many features that could be implemented, especially with respect to pre-built models and tasks. So far I've focused more on the overall structure, than on feature-completeness. Anyone can suggest a feature they'd like to see.
 
-!!! Note:
-    For comments on the documentation, I have specifically enabled to the Giscus commenting system so that GitHub users can comment on pages directly. You can also participate via the Discussions tab on GitHub.
+You're also welcome to leave feedback on the documentation. GitHub users can comment at the bottom of each page of documentation, or directly on [GitHub](https://github.com/mlprt/feedbax/discussions/categories/documentation).
 
 ## Acknowledgments
 
-Special thanks to [Patrick Kidger](https://github.com/patrick-kidger), whose JAX libraries and their documentation often serve as examples to me.
+- Thanks to my PhD supervisor Gunnar Blohm and to the rest of our [lab](http://compneurosci.com/), as well as to Dominik Endres and Stephen H. Scott for discussions that have influenced this project
+- Special thanks to [Patrick Kidger](https://github.com/patrick-kidger), whose JAX libraries and their documentation often serve as examples to me
 

examples/0_train_simple.ipynb

@@ -152,7 +152,7 @@
    "source": [
     "## Interpreting the training history\n",
     "\n",
-    "Feedbax provides some convenient plotting functions, such as [`plot_losses`][feedbax.plot.plot_losses]."
+    "Feedbax provides some convenient plotting functions, such as [`plot_loss_history`][feedbax.plot.plot_loss_history]."
    ]
   },
   {