Repository: ajthor/socks (language: python)

Issue:
separating_kernel.py
**To be double-checked:**
https://github.com/ajthor/socks/blob/47c3cad0c841a530898e2664fc7393b54e70aeab/gym_socks/algorithms/reach/separating_kernel.py#L98 should be replaced with `self.tau = 1 - np.min(np.diagonal(K.T @ self.W @ K))`.

https://github.com/ajthor/socks/blob/47c3cad0c841a530898e2664fc7393b54e70aeab/gym_socks/algorithms/reach/separating_kernel.py#L118 should be replaced with `C = np.diagonal(K.T @ self.W @ K)`.


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 557a752..6fd28d5 100644
--- a/README.md
+++ b/README.md
@@ -40,6 +40,6 @@ In order to cite the toolbox, use the following bibtex entry:
   title={{SOCKS}: A Kernel-Based Stochastic Optimal Control and Reachability Toolbox},
   authors={Thorpe, Adam J. and Oishi, Meeko M. K.},
   year={2022},
-  booktitle={Proceedings of the 25th ACM International Conference on Hybrid Systems: Computation and Control (to appear)},
+  booktitle={Proceedings of the 25th ACM International Conference on Hybrid Systems: Computation and Control},
 }
 ```
diff --git a/benchmarks/kernel_sr/baseline.py b/benchmarks/kernel_sr/baseline.py
index 8369067..a0b2748 100644
--- a/benchmarks/kernel_sr/baseline.py
+++ b/benchmarks/kernel_sr/baseline.py
@@ -17,7 +17,7 @@
 
 from gym_socks.algorithms.reach import kernel_sr
 
-from sklearn.metrics.pairwise import rbf_kernel
+from gym_socks.kernel.metrics import rbf_kernel
 
 
 ex = Experiment()
@@ -139,7 +139,7 @@ def main(
         target_tube=target_tube,
         problem=problem,
         regularization_param=regularization_param,
-        kernel_fn=partial(rbf_kernel, gamma=1 / (2 * (sigma ** 2))),
+        kernel_fn=partial(rbf_kernel, sigma=sigma),
         verbose=False,
     )
 
diff --git a/benchmarks/kernel_sr_max/baseline.py b/benchmarks/kernel_sr_max/baseline.py
index 70d39f5..c40338c 100644
--- a/benchmarks/kernel_sr_max/baseline.py
+++ b/benchmarks/kernel_sr_max/baseline.py
@@ -15,11 +15,9 @@
 from gym_socks.sampling import sample
 from gym_socks.sampling import default_sampler
 
-from gym_socks.utils.grid import make_grid_from_ranges
-
 from gym_socks.algorithms.reach import kernel_sr_max
 
-from sklearn.metrics.pairwise import rbf_kernel
+from gym_socks.kernel.metrics import rbf_kernel
 
 
 ex = Experiment()
@@ -166,7 +164,7 @@ def main(
         target_tube=target_tube,
         problem=problem,
         regularization_param=regularization_param,
-        kernel_fn=partial(rbf_kernel, gamma=1 / (2 * (sigma ** 2))),
+        kernel_fn=partial(rbf_kernel, sigma=sigma),
         verbose=False,
     )
 
diff --git a/docs/bibliography.bib b/docs/bibliography.bib
index 2979921..053056e 100644
--- a/docs/bibliography.bib
+++ b/docs/bibliography.bib
@@ -2,7 +2,7 @@ @inproceedings{thorpe2022hscc
   title     = {{SOCKS}: A Kernel-Based Stochastic Optimal Control and Reachability Toolbox},
   authors   = {Thorpe, Adam J. and Oishi, Meeko M. K.},
   year      = {2022},
-  booktitle = {Proceedings of the 25th ACM International Conference on Hybrid Systems: Computation and Control (to appear)}
+  booktitle = {Proceedings of the 25th ACM International Conference on Hybrid Systems: Computation and Control}
 }
 
 @article{thorpe2020model,
diff --git a/docs/conf.py b/docs/conf.py
index 7faf4e5..8914308 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -105,7 +105,10 @@
 ]
 
 # bibtex
-bibtex_bibfiles = ["bibliography.bib"]
+bibtex_bibfiles = [
+    "bibliography.bib",
+    "guide/kernel_tutorial/bibliography-embeddings.bib",
+]
 
 # Code block and signature options.
 copybutton_prompt_text = r">>> |\.\.\. "
diff --git a/docs/contributing/new_benchmarks.rst b/docs/contributing/new_benchmarks.rst
index 937f034..3cdf964 100644
--- a/docs/contributing/new_benchmarks.rst
+++ b/docs/contributing/new_benchmarks.rst
@@ -14,6 +14,15 @@ published techniques and that show the performance of the algorithms (e.g. scala
 
     Have an idea? Please reach out. We'd love to hear from you!
 
+
+Have a comparison?
+==================
+
+If you have a comparison that you would like to include, please get in touch via the
+`discussion <https://github.com/ajthor/socks/discussions>`_ board. We'd love to help
+include this in the SOCKS toolbox.
+
+
 Current Plan
 ============
 
@@ -231,10 +240,152 @@ slightly more difficult to configure and access.
             if len(results) >= 3:
                 return results
 
+Making Experiments (the old way)
+================================
 
-Have a comparison?
-==================
+We strongly encourage using the `sacred <https://github.com/IDSIA/sacred>`_ experimental
+framework for new experiments. This promotes repeatability and easy configuration that
+is useful for rapid development, iteration, and modification.
 
-If you have a comparison that you would like to include, please get in touch via the
-`discussion <https://github.com/ajthor/socks/discussions>`_ board. We'd love to help
-include this in the SOCKS toolbox.
\ No newline at end of file
+Minimal Experiment
+------------------
+
+The following template is a (mostly) minimal example of how to organize a new
+experiment, and allows for simple modification and extension.
+
+.. code-block:: python
+
+    import gym
+    import gym_socks
+    import logging
+
+    import numpy as np
+
+    from sacred import Experiment
+
+    from examples._computation_timer import ComputationTimer
+
+    from examples.ingredients.system_ingredient import system_ingredient
+    from examples.ingredients.system_ingredient import set_system_seed
+    from examples.ingredients.system_ingredient import make_system
+
+
+    @system_ingredient.config
+    def system_config():
+        system_id = "2DIntegratorEnv-v0"
+        sampling_time = 0.25
+
+
+    ex = Experiment(ingredients=[system_ingredient])
+
+
+    @ex.config
+    def config():
+        time_horizon = 5
+        results_filename = "results/data.npy"
+
+
+    @ex.main
+    def main(seed, time_horizon, results_filename, _log):
+
+        env = make_system()
+
+        # Set the random seed.
+        set_system_seed(seed, env)
+
+        with ComputationTimer():
+
+            # Run algorithm here.
+
+        # Save the result to NPY file.
+        _log.debug(f"Saving the results to file {results_filename}.")
+        with open(results_filename, "wb") as f:
+            np.save(f, result)
+
+
+    if __name__ == "__main__":
+        ex.run_commandline()
+
+The "magic" of sacred comes from the configuration functions, and how the variables
+defined within them can be updated and injected into the ``main`` function. However, the
+main ideas are to define parameters and configuration variables within the ``config``
+function, and to add them as arguments to the ``main`` function. Then, at runtime,
+sacred will pass either the default value specified in the config, or use the updated
+values passed by the user at the command line. In fact, this works for any "captured"
+function, which means that only the relevant configuration variables need to be passed
+to the functions that use them. This allows for more modularized experiments and also
+allows for more convenient interaction by the user.
+
+Adding Plotting
+---------------
+
+Plotting results is typically desired for many experiments, and the sacred framework can
+also be used to create configurable plotting commands.
+
+The following code can be added to the experiment, and then a plot can be generated
+using the command:
+
+.. code-block:: shell
+
+    python -m <experiment> plot_results
+
+.. code-block:: python
+
+    from examples.ingredients.plotting_ingredient import plotting_ingredient
+    from examples.ingredients.plotting_ingredient import update_rc_params
+
+    @plotting_ingredient.config_hook
+    def _plot_config(config, command_name, logger):
+        if command_name in {"main", "plot_results"}:
+            return {
+                "axes": {
+                    "xlabel": r"$x_1$",
+                    "ylabel": r"$x_2$",
+                }
+            }
+
+
+    @ex.command(unobserved=True)
+    def plot_results(plot_cfg, results_filename):
+
+        logging.getLogger("matplotlib").setLevel(logging.WARNING)
+        logging.getLogger("PIL").setLevel(logging.WARNING)
+
+        # Dynamically load for speed.
+        import matplotlib
+
+        matplotlib.use("Agg")
+        update_rc_params(matplotlib, plot_cfg["rc_params"])
+
+        import matplotlib.pyplot as plt
+
+        # Load the result from NPY file.
+        with open(results_filename, "rb") as f:
+            result = np.load(f)
+
+        fig = plt.figure()
+        ax = plt.axes(**plot_cfg["axes"])
+
+        # Plotting code here.
+
+        plt.savefig(plot_cfg["plot_filename"])
+
+The utility of this approach is when the format of the results needs to be changed to
+fit a particular publication, but re-running the experiment can be time-intensive. By
+saving the main algorithm results to a file, and then loading them separately in the
+plotting function, we save the time of having to re-compute the algorithm for plot
+manipulation. The :py:mod:`~examples.ingredients.plotting_ingredient` module implements
+a small ingredient that can be used to modify the ``rc_params`` of ``matplotlib``, or
+add configuration options that can be used to modify the plot appearance. For example,
+
+.. code-block:: shell
+
+    python -m <experiment> plot_results with <updates>
+
+.. hint::
+
+    The plotting config is typically hidden when using the ``print_config`` command
+    using the above code. This is due to the ``if command_name in {"main",
+    "plot_results"}`` line above. To show the config, either use the ``-p`` flag when
+    calling the ``plot_results`` command, or remove the ``if`` statement in the
+    ``_plot_config`` function above.
\ No newline at end of file
diff --git a/docs/guide/for_gym_users.rst b/docs/guide/for_gym_users.rst
new file mode 100644
index 0000000..aee8d12
--- /dev/null
+++ b/docs/guide/for_gym_users.rst
@@ -0,0 +1,129 @@
+********************
+For OpenAI Gym Users
+********************
+
+SOCKS and `OpenAI gym <https://github.com/openai/gym>`_ have a lot in common, but SOCKS
+has some key differences which are different to how gym does things.
+
+In this guide, we go over some of these differences and help get you started with SOCKS.
+
+1. ``gym.Env`` vs. :py:class:`~gym_socks.envs.dynamical_system.DynamicalSystem`
+===============================================================================
+
+An environment in gym (:py:obj:`gym.Env`) is a class with a ``step()``, ``reset()``,
+``render()``, and ``close()`` method. A
+:py:class:`~gym_socks.envs.dynamical_system.DynamicalSystem` inherits from
+:py:obj:`gym.Env`, which means it has all of the same features.
+
+However, gym environments are designed to encapsulate every part of the runtime loop,
+including computing the cost, determining simulation termination, and tracking and
+simulating external "world" objects such as obstacles. This can decrease transparency to
+the designer, and limits modularity in the simulation design process.
+
+In SOCKS, this functionality is handled externally from the environment class.
+Calculating cost, simulating obstacles, determining the conditions for ending
+simulations, etc. are left to the designer.
+
+In addition, there are a few changes to the class interface and the available
+attributes.
+
+:py:meth:`~gym_socks.envs.dynamical_system.DynamicalSystem.dynamics`
+--------------------------------------------------------------------
+
+The :py:meth:`~gym_socks.envs.dynamical_system.DynamicalSystem.step` function for
+dynamical systems generally does not need to be overloaded when defining custom
+environments. The critical component of a dynamical system is the system dynamics. We
+define a new method,
+:py:meth:`~gym_socks.envs.dynamical_system.DynamicalSystem.dynamics`, that is used
+internally by an initial value problem (IVP) solver in the
+:py:meth:`~gym_socks.envs.dynamical_system.DynamicalSystem.step` function. This means we
+only need to define a custom
+:py:meth:`~gym_socks.envs.dynamical_system.DynamicalSystem.dynamics` method when
+subclassing a :py:class:`~gym_socks.envs.dynamical_system.DynamicalSystem`.
+
+:py:attr:`~gym_socks.envs.dynamical_system.DynamicalSystem.state_space`
+-----------------------------------------------------------------------
+
+In gym, an environment nominally has only an
+:py:attr:`~gym_socks.envs.dynamical_system.DynamicalSystem.action_space` and an
+:py:attr:`~gym_socks.envs.dynamical_system.DynamicalSystem.observation_space`. The
+reason for omitting a state space is unclear, since this is an essential component of a
+POMDP definition, which is generally what the :py:obj:`gym.Env` model is based on.
+
+In SOCKS, we correct this by adding a
+:py:attr:`~gym_socks.envs.dynamical_system.DynamicalSystem.state_space` attribute to the
+class. The :py:attr:`~gym_socks.envs.dynamical_system.DynamicalSystem.action_space`
+remains unchanged, but the actual underlying state of the system is an element of the
+:py:attr:`~gym_socks.envs.dynamical_system.DynamicalSystem.state_space`, and the output
+of the :py:meth:`~gym_socks.envs.dynamical_system.DynamicalSystem.step` function is an
+element of the
+:py:attr:`~gym_socks.envs.dynamical_system.DynamicalSystem.observation_space`.
+
+:py:meth:`~gym_socks.envs.dynamical_system.DynamicalSystem.generate_observation`
+--------------------------------------------------------------------------------
+
+Using the underlying mathematical theory, an observation is generated by an *observation
+function*, that takes the current state, action, and a random noise variable
+representing measurement noise (which is different from process noise), and outputs an
+observation. In the :py:class:`~gym_socks.envs.dynamical_system.DynamicalSystem` class,
+we separate this out into its own method.
+
+Classes inheriting from :py:class:`~gym_socks.envs.dynamical_system.DynamicalSystem`
+should override this method in order to include measurement noise or limit the state
+variables that are observable. By default, this function simply returns the true
+underlying state of the system (meaning the system is fully observable).
+
+:py:meth:`~gym_socks.envs.dynamical_system.DynamicalSystem.generate_disturbance`
+--------------------------------------------------------------------------------
+
+We also separate the disturbance from the
+:py:meth:`~gym_socks.envs.dynamical_system.DynamicalSystem.step` and
+:py:meth:`~gym_socks.envs.dynamical_system.DynamicalSystem.dynamics` functions. In
+SOCKS, we generate the disturbance independently in its own method.
+
+Classes inheriting from :py:class:`~gym_socks.envs.dynamical_system.DynamicalSystem`
+should override this method to specify the type of disturbance which affects the
+dynamical evolution of the system, for instance to increase or decrease the process
+noise that affects the state or change the noise distribution.
+
+2. Policies
+===========
+
+In SOCKS, we typically represent the policy controlling a system as a separate object or
+function. During simulation, the policy is a function that returns a control action (and
+may depend on the simulation time and/or the system state).
+
+SOCKS defines a class, :py:class:`~gym_socks.policies.policy.BasePolicy`, that has a
+simple ``Callable`` interface, meaning it requires that the policy implement a
+:py:meth:`~gym_socks.policies.policy.BasePolicy.__call__` method.
+
+This class is used by some of the algorithms, but is generally useful more for code
+organization and to provide a consistent interface, rather than enforcing a strict
+simulation scheme.
+
+3. Sampling
+===========
+
+A core component of SOCKS is the ability to generate a finite sample (a collection of
+observations) from a system. In gym, the simulation is typically handled via a ``for``
+loop, which also includes the code used to collect observations and train the policy.
+
+In SOCKS, we separate this process into a sampling phase and a separate training phase.
+The sampling phase consists of collecting a set of observations of the system
+transitions, and is handled by the various sampling functions implemented in SOCKS. The
+training phase is then typically handled via the algorithm, which uses the sample to
+compute a policy.
+
+Stable Baselines
+================
+
+The algorithms in `stable-baselines3 <https://github.com/DLR-RM/stable-baselines3>`_ are
+designed to be used specifically in reinforcement learning training loops.
+
+SOCKS, on the other hand, is geared toward solving stochastic optimal control problems
+(it can be argued that RL is a special case of stochastic optimal control). With this in
+mind, we have emphasized modularity in our design, and have abstracted many components
+of the training and simulation loop away from the :py:obj:`gym.Env` class.
+
+If you are already familiar with ``stable-baselines3``, then it may be helpful to think
+of SOCKS as a set of algorithms based in kernel methods for solving similar problems.
\ No newline at end of file
diff --git a/docs/guide/kernel_tutorial/bibliography-embeddings.bib b/docs/guide/kernel_tutorial/bibliography-embeddings.bib
new file mode 100644
index 0000000..0d5387a
--- /dev/null
+++ b/docs/guide/kernel_tutorial/bibliography-embeddings.bib
@@ -0,0 +1,84 @@
+@inproceedings{10.1007/978-3-540-75225-7_5,
+  author    = {Smola, Alex and Gretton, Arthur and Song, Le and Sch\"{o}lkopf, Bernhard},
+  title     = {A {Hilbert} Space Embedding for Distributions},
+  year      = {2007},
+  isbn      = {9783540752240},
+  publisher = {Springer-Verlag},
+  address   = {Berlin, Heidelberg},
+  doi       = {10.1007/978-3-540-75225-7_5},
+  abstract  = {We describe a technique for comparing distributions without the need for density estimation as an intermediate step. Our approach relies on mapping the distributions into a reproducing kernel Hilbert space. Applications of this technique can be found in two-sample tests, which are used for determining whether two sets of observations arise from the same distribution, covariate shift correction, local learning, measures of independence, and density estimation.},
+  booktitle = {Proceedings of the 18th International Conference on Algorithmic Learning Theory},
+  pages     = {13–31},
+  numpages  = {19},
+  location  = {Sendai, Japan},
+  series    = {ALT '07}
+}
+
+@inproceedings{10.1145/1553374.1553497,
+  author    = {Song, Le and Huang, Jonathan and Smola, Alex and Fukumizu, Kenji},
+  title     = {{Hilbert} Space Embeddings of Conditional Distributions with Applications to Dynamical Systems},
+  year      = {2009},
+  isbn      = {9781605585161},
+  publisher = {Association for Computing Machinery},
+  address   = {New York, NY, USA},
+  doi       = {10.1145/1553374.1553497},
+  abstract  = {In this paper, we extend the Hilbert space embedding approach to handle conditional distributions. We derive a kernel estimate for the conditional embedding, and show its connection to ordinary embeddings. Conditional embeddings largely extend our ability to manipulate distributions in Hilbert spaces, and as an example, we derive a nonparametric method for modeling dynamical systems where the belief state of the system is maintained as a conditional embedding. Our method is very general in terms of both the domains and the types of distributions that it can handle, and we demonstrate the effectiveness of our method in various dynamical systems. We expect that conditional embeddings will have wider applications beyond modeling dynamical systems.},
+  booktitle = {Proceedings of the 26th Annual International Conference on Machine Learning},
+  pages     = {961–968},
+  numpages  = {8},
+  location  = {Montreal, Quebec, Canada},
+  series    = {ICML '09}
+}
+
+@article{MAL-060,
+  year    = {2017},
+  volume  = {10},
+  journal = {Foundations and Trends® in Machine Learning},
+  title   = {Kernel Mean Embedding of Distributions: A Review and Beyond},
+  doi     = {10.1561/2200000060},
+  issn    = {1935-8237},
+  number  = {1-2},
+  pages   = {1-141},
+  author  = {Krikamol Muandet and Kenji Fukumizu and Bharath Sriperumbudur and Bernhard Schölkopf}
+}
+
+@inproceedings{10.5555/3042573.3042778,
+  author    = {Gr\"{u}new\"{a}lder, Steffen and Lever, Guy and Baldassarre, Luca and Pontil, Massimilano and Gretton, Arthur},
+  title     = {Modelling Transition Dynamics in MDPs with RKHS Embeddings},
+  year      = {2012},
+  isbn      = {9781450312851},
+  publisher = {Omnipress},
+  address   = {Madison, WI, USA},
+  abstract  = {We propose a new, nonparametric approach to learning and representing transition dynamics in Markov decision processes (MDPs), which can be combined easily with dynamic programming methods for policy optimisation and value estimation. This approach makes use of a recently developed representation of conditional distributions as embeddings in a reproducing kernel Hilbert space (RKHS). Such representations bypass the need for estimating transition probabilities or densities, and apply to any domain on which kernels can be defined. This avoids the need to calculate intractable integrals, since expectations are represented as RKHS inner products whose computation has linear complexity in the number of points used to represent the embedding. We provide guarantees for the proposed applications in MDPs: in the context of a value iteration algorithm, we prove convergence to either the optimal policy, or to the closest projection of the optimal policy in our model class (an RKHS), under reasonable assumptions. In experiments, we investigate a learning task in a typical classical control setting (the under-actuated pendulum), and on a navigation problem where only images from a sensor are observed. For policy optimisation we compare with least-squares policy iteration where a Gaussian process is used for value function estimation. For value estimation we also compare to the NPDP method. Our approach achieves better performance in all experiments.},
+  booktitle = {Proceedings of the 29th International Coference on Machine Learning},
+  pages     = {1603–1610},
+  numpages  = {8},
+  location  = {Edinburgh, Scotland},
+  series    = {ICML'12}
+}
+
+@article{6530747,
+  author  = {Song, Le and Fukumizu, Kenji and Gretton, Arthur},
+  journal = {IEEE Signal Processing Magazine},
+  title   = {Kernel Embeddings of Conditional Distributions: A Unified Kernel Framework for Nonparametric Inference in Graphical Models},
+  year    = {2013},
+  volume  = {30},
+  number  = {4},
+  pages   = {98-111},
+  doi     = {10.1109/MSP.2013.2252713}
+}
+
+@inproceedings{10.5555/3042573.3042803,
+  author    = {Gr\"{u}new\"{a}lder, Steffen and Lever, Guy and Baldassarre, Luca and Patterson, Sam and Gretton, Arthur and Pontil, Massimilano},
+  title     = {Conditional Mean Embeddings as Regressors},
+  year      = {2012},
+  isbn      = {9781450312851},
+  publisher = {Omnipress},
+  address   = {Madison, WI, USA},
+  abstract  = {We demonstrate an equivalence between reproducing kernel Hilbert space (RKHS) embeddings of conditional distributions and vector-valued regressors. This connection introduces a natural regularized loss function which the RKHS embeddings minimise, providing an intuitive understanding of the embeddings and a justification for their use. Furthermore, the equivalence allows the application of vector-valued regression methods and results to the problem of learning conditional distributions. Using this link we derive a sparse version of the embedding by considering alternative formulations. Further, by applying convergence results for vector-valued regression to the embedding problem we derive minimax convergence rates which are O(log(n)/n) - compared to current state of the art rates of O(n-1/4) - and are valid under milder and more intuitive assumptions. These minimax upper rates coincide with lower rates up to a logarithmic factor, showing that the embedding method achieves nearly optimal rates. We study our sparse embedding algorithm in a reinforcement learning task where the algorithm shows significant improvement in sparsity over an incomplete Cholesky decomposition.},
+  booktitle = {Proceedings of the 29th International Conference on Machine Learning},
+  pages     = {1803–1810},
+  numpages  = {8},
+  location  = {Edinburgh, Scotland},
+  series    = {ICML'12}
+}
\ No newline at end of file
diff --git a/docs/guide/kernel_tutorial/nn_comparison.rst b/docs/guide/kernel_tutorial/nn_comparison.rst
new file mode 100644
index 0000000..11ba592
--- /dev/null
+++ b/docs/guide/kernel_tutorial/nn_comparison.rst
@@ -0,0 +1,197 @@
+**************************
+Neural Network vs. Kernels
+**************************
+
+Several recent papers have begun to dive into the implicit connections between neural
+networks and kernel methods. They claim that understanding kernel methods will be
+critical in order to understand and overcome some of the theoretical roadblocks that
+have cropped up at the forefront of neural network research.
+
+.. note::
+
+    The connection between these two competing techniques has been studied extensively, so here we describe only the basic case of a single-layer neural network.
+
+Feature Maps
+============
+
+You may have heard the term "feature" in learning literature, as in, machine learning
+attempts to learn data *features* (which can potentially be described as learning the
+underlying patterns in the data).
+
+In practice, this often amounts to learning whether the data fits a set of pre-defined
+functions :math:`\phi_{i}`, :math:`i = 1, 2, \ldots`, and we combine these features
+using some sort of linear combination,
+
+.. math::
+
+    f(x) = \sum_{i=1}^{M} \alpha_{i} \phi_{i}(x)
+
+In kernel methods, we have a very simple way of defining features via the kernel. In
+effect, the features are simply the kernel function fixed at a number of pre-specified
+points, i.e. :math:`\phi_{i}(\cdot) = k(x_{i}, \cdot)`.
+
+.. math::
+
+    f(x) = \sum_{i=1}^{M} \alpha_{i} k(x_{i}, x)
+
+There are a number of well-known kernels that have been found, and are typically used in
+kernel methods, such as the constant kernel :math:`k(x, x') = a`, linear kernel
+:math:`k(x, x') = \langle x, x' \rangle`, polynomial kernel :math:`k(x, x') = (\langle
+x, x' \rangle + c)^{d}`, and the Gaussian (RBF) kernel :math:`k(x, x') = \exp(-\lVert x
+- x' \rVert^{2}/2 \sigma^{2})`.
+
+We can also combine these kernels as needed, for example by adding them together,
+scaling them, or multiplying two kernels, in order to create new valid kernels. This
+allows us to define the "features" easily via the kernel.
+
+.. hint::
+
+    What this also means is, a kernel function can itself be a linear combination of
+    kernels.
+
+Approximation via Features
+--------------------------
+
+We can describe the learning problem in terms of features, meaning for a set of
+pre-defined features :math:`\lbrace \phi_{i} \rbrace_{i=1}^{M}`, we seek to learn the
+coefficients :math:`\alpha_{i}` of a linear combination of feature maps,
+
+.. math::
+
+    f(x) = \alpha_{1} \phi_{1}(x) + \cdots + \alpha_{M} \phi_{M}(x)
+
+This is exactly the same as the kernel learning problem that we have described before.
+In short, whether we call the kernels "features" or vice versa, we are essentially using
+the same techniques that we have already described in the first part of the
+:doc:`tutorial </guide/kernel_tutorial/tutorial>`.
+
+However, it is important to note that **not all feature maps are valid kernels**. Some
+features are complex, especially in the case of deep neural networks or convolutional
+networks, which means not all neural networks are directly relatable to kernel methods.
+Nevertheless, in certain cases, these two techniques are virtually identical.
+
+Single-Layer Neural Networks
+============================
+
+In a single-layer neural network, we take an input vector :math:`x`, and compute a
+linear transformation of the data, :math:`W x + b`, where :math:`W` is usually described
+as a matrix of "weights", and :math:`b` is called a "bias" vector.
+
+We then pass the result through a nonlinear transformation :math:`g`, which is called
+the "activation" function. At the final step, we may then compute some linear
+combination of the "activations" in order to compute the predicted value of the function
+approximation.
+
+.. math::
+
+    f(x) = \sum_{i=1}^{M} \beta_{i} g_{i}(W x + b)
+
+We specify the activation functions explicitly, and the neural network learning problem
+is defined as learning the weights :math:`W`, the bias :math:`b`, and the linear
+combination coefficients :math:`\beta_{i}`. Thanks to a technique called
+*backpropagation*, we can directly compute the gradient of these values with respect to
+the error. This means we can use an iterative training scheme called gradient descent,
+where the gradient is used to update the weights, biases, and coefficients of the neural
+network in order to iteratively improve the neural network approximation.
+
+However, it is important to note that this procedure does not guarantee that we will
+converge to a perfect solution. In practice, if we do not take sufficiently small steps
+when we update the neural network parameters, we may fail to minimize the error of the
+neural network approximator.
+
+In order to describe a single-layer neural network in terms of features, we define the
+feature map to be the activations of the neural network.
+
+.. math::
+
+    \varphi_{i}(x) = g_{i}(W x + b)
+
+Then, it is easy to see that the neural network follows the same scheme as the kernel
+learning problem when viewed from the last layer.
+
+.. math::
+
+    f(x) = \beta_{1} \varphi_{1}(x) + \cdots + \beta_{M} \varphi_{M}(x)
+
+This is an important insight, since it gives us a way to relate the kernel features
+:math:`\phi_{i}` and the neural network features :math:`\varphi_{i}`. In certain cases,
+these feature maps may be directly related, meaning what we are doing in neural network
+learning problems is effectively "learning the kernel".
+
+.. note::
+
+    Note that deep neural networks and convolutional networks are different, and the
+    feature maps for these networks are significantly more complex. These networks
+    typically do not have a simple analog to kernel methods.
+
+Practical Differences
+=====================
+
+There are also some practical differences between neural networks and kernel methods
+that we outline below.
+
+Input Dimensionality
+--------------------
+
+Note that the dimensionality of the input vector plays a role in the complexity of a
+neural network. If the dimensionality of the input vector :math:`x` is large, the weight
+matrix :math:`W` will also have very high dimensionality. This can make maintaining and
+updating the weights challenging, since we need to explicitly store and compute the
+gradient updates for a high-dimensional matrix.
+
+Kernel methods do not suffer from this problem, since the data is projected into a
+high-dimensional function space, and we can use the :doc:`kernel trick
+</guide/kernel_tutorial/rkhs_tutorial>`. This means we only need
+to evaluate the *kernel* over high-dimensional data, which is typically much easier
+since it boils down to an inner product of vectors.
+
+Sequential Data
+---------------
+
+Neural networks are particularly adept at handling data which is provided
+*sequentially*, meaning we receive observations given one at a time from the true
+function. Since we can update the parameters of the network as data becomes available,
+we do not need to store all of the data that has been passed previously.
+
+The solution for kernel methods, however, requires that we store a matrix that scales
+with the amount of data. This is a major hurdle for kernel methods, since practically it
+requires us to first collect a large sample, which we can then use to easily compute the
+"best fit" solution.
+
+In short, what this means is that we are presented with a tradeoff. Neural networks take
+a long time to train, but are relatively cheap to store in memory once the parameters
+have been learned. Kernel methods, on the other hand, face large memory storage
+requirements, but are trained in a single calculation. Despite this, neural networks are
+typically better calibrated for sequential data scenarios when a large amount of data is
+involved.
+
+.. seealso::
+
+    Note that several techniques for reducing the computational complexity of kernel
+    methods have been developed, such as Random Fourier Features, the Nystrom Method,
+    and Gaussian matrix approximations. Check out the :doc:`/examples/index` page for
+    a demonstrations of how these techniques work.
+
+Theoretical Guarantees
+----------------------
+
+One major disadvantage of neural networks is the lack of theoretical guarantees related
+to convergence of the solution and what is called "explainability", i.e. why the network
+has a particular output in response to a given input. This is a deep topic, so we will
+only cover it here briefly.
+
+Basically, the training procedure for neural networks is not conducive to
+"introspection". Tracking an output backwards through the network leads to a complex web
+of interconnected weights and activations, so determining the reasons why a network may
+output a poor prediction is an area of open research.
+
+In addition, the training for neural networks is typically iterative, and the solution
+can jump around or get caught in a local solution without ever converging to the true,
+global solution. Mathematically, this is because the neural network learning problem is
+*non-convex*.
+
+Kernel learning, on the other hand, is a *convex* problem, meaning kernel-based
+solutions do not suffer from local minima, and are guaranteed to find the "best fit" for
+a given data set. In addition, kernel methods are typically "shallow", meaning they
+offer better explainability in their outputs. In other words, we can always find which
+features and training data contributed to a certain prediction.
diff --git a/docs/guide/kernel_tutorial/rkhs_tutorial.rst b/docs/guide/kernel_tutorial/rkhs_tutorial.rst
new file mode 100644
index 0000000..b274325
--- /dev/null
+++ b/docs/guide/kernel_tutorial/rkhs_tutorial.rst
@@ -0,0 +1,160 @@
+*************
+RKHS Tutorial
+*************
+
+In this tutorial, we dive a little bit deeper into the theory of a reproducing kernel
+Hilbert space and give a brief overview of one of the main properties of an RKHS, known
+as the *reproducing property*.
+
+*So what is an RKHS, really?*
+=============================
+
+In the first part of the :doc:`tutorial </guide/kernel_tutorial/tutorial>`, we described
+the main idea of a kernel and kernel methods. In this part of the tutorial, we are going
+to give a brief overview of an RKHS, and explain what that means in basic terms.
+
+.. admonition:: Did you know?
+
+    It might surprise you to know that you have likely been working with Hilbert spaces
+    all along without even knowing it. The most common Hilbert space is
+    :math:`\mathbb{R}^{n}`, which means that if you've taken any high school math course
+    or linear algebra, you've probably already had a basic introduction to Hilbert
+    spaces.
+
+To recap, we started with a kernel function :math:`k : \mathcal{X} \times \mathcal{X}
+\to \mathbb{R}`, and we said that by fixing one of the parameters, we get a kernel that
+is "centered" at a particular value. In other words, for some fixed :math:`x`, we get a
+function, :math:`k(x, \cdot)`. If we do this for all possible values of :math:`x`, what
+we get is a set of functions :math:`\lbrace k(x, \cdot) \mid x \in \mathcal{X} \rbrace`.
+The set of all of these functions is what we are going to use to construct our RKHS.
+However, depending on the space :math:`\mathcal{X}`, it is easy to see that there might
+be infinitely many functions in the set (for example, if :math:`\mathcal{X} =
+\mathbb{R}`).
+
+Before we continue, we are going to review the basic mathematical idea of vector spaces.
+A vector space is a set of mathematical objects that can be added together and
+multiplied by a scalar value. The elements of a vector space are called *vectors*. The
+important point here is that *any space that satisfies the properties of a vector space
+can be called a vector space*, and the elements of this space are called vectors.
+
+Looking ahead, what we would like to do is combine these functions and multiply them by
+some scalar value in order to compute a function approximation,
+
+.. math::
+    :label: linear_combination
+
+    f(x) = \sum_{i=1}^{\infty} \alpha_{i} k(x_{i}, x)
+
+This means if we can add our kernel functions :math:`k(x, \cdot)` and scale them by some
+scalar value :math:`\alpha`, then the functions :math:`k(x, \cdot)` are vectors.
+
+In essence, this is exactly what we have in :eq:`linear_combination`.
+
+What this means is that the RKHS is a vector space, but instead of an array of numbers,
+we have *functions*. Nevertheless, all of the properties of vector spaces still apply.
+
+Inner Products
+--------------
+
+Now, we can add two vectors (kernel functions) together and multiply them by scalars,
+but we cannot "multiply" two vectors together yet. To do this, we define what is known
+as the "inner" product. When we're dealing with the familiar concept of vectors in
+:math:`\mathbb{R}^{n}`, we have a special operation we can perform called the "dot"
+product, :math:`x^{\top} y`, :math:`x, y \in \mathbb{R}^{n}`.
+
+Mathematically, this is simply a different way of writing the inner product between
+vectors, which is typically denoted as:
+
+.. math::
+
+    \langle x, y \rangle, \quad x, y \in \mathbb{R}^{n}
+
+In order to turn our set of functions :math:`\lbrace k(x, \cdot) \mid x \in \mathcal{X}
+\rbrace` into a Hilbert space, we need to define an inner product. Basically, we need to
+explicitly define what it means to "multiply" two vectors together. Luckily, this is
+already defined for us by the kernel, which is what makes kernels so special.
+
+When we take one element from our set :math:`k(x, \cdot)`, and we take the inner product
+with another element from our set :math:`k(x', \cdot)`, what we get is,
+
+.. math::
+    :label: kernel_trick
+
+    \langle k(x, \cdot), k(x', \cdot) \rangle = k(x, x')
+
+This has another name in kernel methods, and is called the *"kernel trick"*.
+
+.. note::
+
+    Once we have defined an inner product, there are a few more mathematical steps to
+    define an RKHS. For the interested reader, we need to ensure that the space is also
+    complete, meaning the space contains all of its limit points. Mathematically, this
+    means an RKHS :math:`\mathscr{H}` is defined as the closure of the span of kernel
+    functions,
+
+    .. math::
+
+        \mathscr{H} = \overline{\mathrm{span} \lbrace k(x, \cdot) \mid x \in
+        \mathcal{X} \rbrace}
+
+Functions in an RKHS
+--------------------
+
+To recap, what this means is that we start with a kernel function, and we use that
+kernel function to build a set :math:`\lbrace k(x, \cdot) \mid x \in \mathcal{X}
+\rbrace`. The kernel function also gives us an *inner product* via :eq:`kernel_trick`.
+So now we can add the functions in our set, scale them, and take the inner product.
+
+Then, the functions in an RKHS are relatively straightforward to describe. In effect,
+they are some linear combination of kernel functions, just like we saw in the previous
+:doc:`tutorial </guide/kernel_tutorial/tutorial>`. What this means is that any function
+in the RKHS looks like a sum of weighted kernel functions,
+
+.. math::
+    :label: rkhs_function
+
+    f(\cdot) = \sum_{i=1}^{\infty} \alpha_{i} k(x_{i}, \cdot)
+
+where the coefficients :math:`\alpha_{i}` are what differentiate them from each other.
+
+The Reproducing Property
+------------------------
+
+This leads us to the final property of an RKHS that is perhaps the most critical. This
+property is known as the **"reproducing property"**. In a nutshell, the reproducing
+property tells us that we can use the inner product to evaluate a function in the RKHS
+at a particular point. Mathematically, it looks like this:
+
+.. math::
+
+    f(x) = \langle f, k(x, \cdot) \rangle
+
+In order to gain a better understanding of how this works, let's substitute in the
+definition of a function in the RKHS from :eq:`rkhs_function` and use the inner product
+in :eq:`kernel_trick`.
+
+.. math::
+
+    f(x) = \langle f, k(x, \cdot) \rangle
+    = \biggl\langle \sum_{i=1}^{\infty} \alpha_{i} k(x_{i}, \cdot), k(x, \cdot)
+    \biggr\rangle = \sum_{i=1}^{\infty} \alpha_{i} \langle k(x_{i}, \cdot), k(x, \cdot)
+    \rangle = \sum_{i=1}^{\infty} \alpha_{i} k(x_{i}, x)
+
+Now, you may be wondering how we can take the infinite sum of kernel functions in order
+to compute the function. The short answer is that generally, we can't. However, not all
+kernels lead to infinite sums, and the kernel trick is a useful tool here. In general,
+though, this is why we have function *approximations*, where instead of taking an
+infinite sum, we approximate the function as a **finite** sum of kernel functions at a
+set of known points (our data).
+
+Recap
+=====
+
+At the end of this tutorial, you should have some basic understanding of kernels and
+reproducing kernel Hilbert spaces. You should be able to describe what the functions
+inside the RKHS look like, and how the reproducing property allows us to evaluate
+functions.
+
+We've merely scratched the surface on reproducing kernel Hilbert spaces, though. While
+the basic theory is fairly straightforward, there is a large amount of research in this
+area that comes directly from learning theory, probability, and functional analysis.
\ No newline at end of file
diff --git a/docs/guide/kernel_tutorial/statistical_learning.py b/docs/guide/kernel_tutorial/statistical_learning.py
new file mode 100644
index 0000000..1de6c05
--- /dev/null
+++ b/docs/guide/kernel_tutorial/statistical_learning.py
@@ -0,0 +1,177 @@
+# %% [markdown]
+"""
+# Statistical Learning
+
+In this tutorial, we will briefly cover some of the techniques that we employ in SOCKS.
+In a previous [tutorial](tutorial.py), we covered the basics of kernels and described
+how kernel methods can be used for function approximation.
+
+"""
+
+# %% nbsphinx="hidden"
+import numpy as np
+from gym_socks.kernel.metrics import rbf_kernel
+from gym_socks.kernel.metrics import regularized_inverse
+
+from matplotlib import pyplot as plt
+
+# %% [markdown]
+# To recap, here is the example from the previous [tutorial](tutorial.py).
+
+# %%
+f = lambda x: x ** 3 - 4 * x ** 2 + 6 * x - 24 + np.exp(-x)
+
+X = np.linspace(-5, 5, 200).reshape(-1, 1)  # Points we want to evaluate at.
+Y = f(X)  # The actual function values. Used for plotting.
+
+sample_size = 100
+x_data = 8 * np.random.rand(sample_size, 1) - 4  # Random data in (-4, 4)
+y_data = f(x_data)
+
+G = rbf_kernel(x_data, sigma=1)
+alpha = np.dot(regularized_inverse(G, regularization_param=1e-9), y_data)
+
+plt.figure()
+plt.axis()
+plt.scatter(x_data, y_data, s=5, c="red", label="Data")
+plt.plot(X, Y, ":k", label="Actual Function")
+plt.plot(X, np.dot(rbf_kernel(X, x_data, sigma=1), alpha), label="Kernel Approximation")
+plt.legend()
+plt.show()
+
+# %% [markdown]
+# Here, we collect data from a function, with points taken randomly in $x$, and we
+# compute the kernel-based approximation as the solution to a regularized least-squares
+# problem. However, note that the data we collect $\lbrace (x_{i}, y_{i})
+# \rbrace_{i=1}^{M}$ is not corrupted by any noise.
+
+# %% [markdown]
+# ## Noisy Measurements
+#
+# What happens when we add some noise to the outputs of our function data?
+
+# %%
+sample_size = 100
+x_data = 8 * np.random.rand(sample_size, 1) - 4  # Random data in (-4, 4)
+# Add Gaussian noise to the data.
+y_data = f(x_data) + 10 * np.random.randn(sample_size, 1)
+
+# Points to evaluate our kernel approximation at.
+x_eval = np.linspace(-5, 5, 200).reshape(-1, 1)
+
+# Compute the kernel-based approximation.
+G = rbf_kernel(x_data, sigma=1)
+y_pred = (
+    y_data.T
+    @ regularized_inverse(G, regularization_param=1e-3)
+    @ rbf_kernel(x_data, x_eval, sigma=1)
+)
+
+plt.figure()
+plt.axis()
+plt.scatter(x_data, y_data, s=5, c="red", label="Data")
+plt.plot(X, Y, ":k", label="Actual Function")
+plt.plot(x_eval, y_pred.T, label="Kernel Approximation")
+plt.legend()
+plt.show()
+
+# %% [markdown]
+# We can see that the function that we learn still fits the noisy data (though not as
+# well), but we need to raise the regularization parameter to keep the solution
+# "smooth". Clearly, the quality of the approximation depends on the data. If we have
+# too much noise in a particular region, or not enough data, the approximation will be
+# poor. We can improve the quality of the approximation by increasing the sample size.
+
+# %%
+sample_size = 500
+x_data = 8 * np.random.rand(sample_size, 1) - 4  # Random data in (-4, 4)
+# Add Gaussian noise to the data.
+y_data = f(x_data) + 10 * np.random.randn(sample_size, 1)
+
+# Points to evaluate our kernel approximation at.
+x_eval = np.linspace(-5, 5, 200).reshape(-1, 1)
+
+# Compute the kernel-based approximation.
+G = rbf_kernel(x_data, sigma=1)
+y_pred = (
+    y_data.T
+    @ regularized_inverse(G, regularization_param=1e-3)
+    @ rbf_kernel(x_data, x_eval, sigma=1)
+)
+
+plt.figure()
+plt.axis()
+plt.scatter(x_data, y_data, s=5, c="red", label="Data")
+plt.plot(X, Y, ":k", label="Actual Function")
+plt.plot(x_eval, y_pred.T, label="Kernel Approximation")
+plt.legend()
+plt.show()
+
+# %% [markdown]
+# Now, it is important to ask, *what exactly are we learning here?*
+#
+# Since the function outputs are corrupted by noise, we are working with a *stochastic
+# function*, meaning it incorporates some randomness. The answer, then, is that we are
+# learning the **expectation** of the stochastic function.
+#
+# In short, if we were to evaluate the stochastic function at the same value of $x$ over
+# and over, we would get different random values for $y$. Thus, if we were simply trying
+# to find the "best fit" function to the random data, we would not actually be fitting
+# the underlying function.
+#
+# Instead, we are attempting to learn an approximation of the function that gives us the
+# *most likely* value of $y$ for every value of $x$. In probability theory, this is
+# called the *expected value* or *expectation*.
+
+# %% [markdown]
+# ## Probability Distributions
+#
+# When we are dealing with random variables and stochasticity, we have what is called a
+# *distribution*, which maps probabilities to outcomes. The distribution also tells us
+# which outcome is the most likely, and this is effectively the expectation.
+#
+# Mathematically, we write this as,
+#
+# $$\mathbb{E}[f(X)]$$
+#
+# Now, if $f$ is a function in an RKHS, we can actually write this using the reproducing
+# property (see [here](rkhs_tutorial.rst)) as the inner product between the function $f$
+# and another element in the RKHS which we call the *kernel distribution embedding*.
+#
+# $$\mathbb{E}[f(X)] = \langle f, m(X) \rangle$$
+#
+# We can compute the expectation at $X = x$ if we fix the random variable $X$ to a
+# particular value $x$.
+#
+# This means that what we are trying to learn is the embedding $m$ in the RKHS.
+
+# %% [raw] raw_mimetype="text/restructuredtext"
+# .. note::
+#
+#     This is a very simple explanation of distributions and expectations, and we have
+#     avoided a more involved explanation to avoid confusion.
+#     Mathematically, these concepts have much more involved definitions.
+#     For those who work with stochasticity or probability in continuous spaces, we
+#     highly recommend checking out a book on measure theory.
+
+# %% [markdown]
+# ### Learning Embeddings
+#
+# Since the embedding $m$ is an element of the RKHS, it can be represented as a linear
+# combination of kernel functions, just as we did before.
+#
+# In practice, this means that the same techniques we used to compute the kernel-based
+# approximation still apply. However, we need to think a little more deeply about what
+# we are actually trying to learn when stochasticity and distributions are involved.
+
+# %% [raw] raw_mimetype="text/restructuredtext"
+# References
+# ----------
+#
+# Kernel embeddings of distributions is a deep topic. Here, we provide a list of some
+# relevant references to get started with kernel embeddings.
+#
+# .. bibliography:: bibliography-embeddings.bib
+#     :list: bullet
+#     :style: unsrt
+#     :all:
diff --git a/docs/guide/kernel_tutorial/tutorial.py b/docs/guide/kernel_tutorial/tutorial.py
new file mode 100644
index 0000000..c5c7fbd
--- /dev/null
+++ b/docs/guide/kernel_tutorial/tutorial.py
@@ -0,0 +1,282 @@
+# %% [markdown]
+"""
+# A Brief Tutorial on Kernels
+
+Kernel methods are a broad class of statistical learning techniques that leverage
+functional analysis and reproducing kernel Hilbert spaces (RKHS) to perform
+classification, clustering, estimation, and regression. Here, we outline the main ideas
+of kernel methods to help build an intuitive understanding of how kernel methods work.
+
+"""
+
+# %% nbsphinx="hidden"
+import numpy as np
+from gym_socks.kernel.metrics import rbf_kernel
+from gym_socks.kernel.metrics import regularized_inverse
+
+from matplotlib import pyplot as plt
+
+# %% [markdown]
+# ## Basics
+
+# %% [markdown]
+# ### *What is a kernel?*
+#
+# Intuitively, a kernel is a function that computes the "similarity" between two things.
+# For instance, given two vectors, a kernel tells us how "similar" the two vectors are.
+# This is almost the same as a distance function, which tells us how "close" two things
+# are--but instead of assigning lower values to objects which are close, a kernel
+# typically assigns a higher value to things that are similar.
+#
+# Mathematically, a kernel is a positive definite function, $k : \mathcal{X} \times
+# \mathcal{X} \to \mathbb{R}$, meaning given a sample $\lbrace x_{i} \rbrace_{i=1}^{M}$,
+# the matrix $G \in \mathbb{R}^{M \times M}$ where $G_{ij} = k(x_{i}, x_{j})$ is
+# positive *semi-*definite.
+#
+# One of the simplest kernels is called the linear kernel, which is defined as the inner
+# product between two things: $k(x, x') = \langle x, x' \rangle$. But perhaps the most
+# common kernel is the Gaussian kernel, also known as the Radial Basis Function (RBF)
+# kernel,
+#
+# $$k(x, x') = \exp\biggl(-\frac{\lVert x - x' \rVert^{2}}{2 \sigma^{2}}\biggr)$$
+#
+# Note that the kernel is a function of two variables. If we fix one of those variables,
+# we effectively center the kernel at a particular value. For instance, with the
+# Gaussian kernel, if we fix one parameter, the Gaussian (mean) becomes centered at that
+# value.
+
+# %%
+X = np.linspace(-5, 5, 200).reshape(-1, 1)  # Points we want to evaluate at.
+
+x_data = np.array([[0]])  # Gaussian center. Data is organized in rows.
+# RBF kernel. Sigma is the "bandwidth" parameter. Larger values give a "wider" kernel.
+k = rbf_kernel(X, x_data, sigma=1)
+
+plt.figure()
+plt.axes()
+plt.plot(X, k)
+plt.show()
+
+# %% [markdown]
+# We can add more Gaussians by defining more "centers".
+
+# %%
+x_data = np.array([[-2], [0], [2]])
+
+plt.figure()
+plt.axes()
+plt.plot(X, rbf_kernel(X, x_data, sigma=1))
+plt.show()
+
+# %% [markdown]
+# The bandwidth $\sigma > 0$ of the Gaussian kernel is known as a "hyperparameter". This
+# is a value we can adjust prior to training, and does not depend explicitly on the
+# data. If we adjust the bandwidth, we get narrower or wider Gaussians.
+
+# %%
+plt.figure()
+plt.axes()
+plt.plot(X, rbf_kernel(X, x_data, sigma=0.5))
+plt.show()
+
+# %% [markdown]
+# ### *What are kernel methods?*
+#
+# Kernel methods refers to a broad class of machine learning and statistical learning
+# techniques that are used in a wide variety of applications. At the core, this
+# generally involves *learning a function* that is represented by kernels.
+#
+# So imagine we have an unknown function $f(x) = y$ that we want to learn, but let's
+# assume that we have some data points $\lbrace (x_{i}, y_{i}) \rbrace_{i=1}^{M}$ that
+# we've observed or that are given to us.
+#
+# The question we want to ask is: *how can we "learn" this function using the data?*
+# What we mean is, we want to find some way to represent the function, and we want to
+# incorporate some tunable parameters or values that we can compute. If we tune the
+# parameters appropriately, the representation we have should closely approximate the
+# true function.
+#
+# For example, we could represent the function by a polynomial, and we use the
+# polynomial coefficients as our tunable parameters. If we represent the function by a
+# line or set of lines, we want to adjust the slope and $y$-intersect of those lines to
+# get a good approximation.
+#
+# This is effectively how all learning works. For instance, neural networks represent a
+# function using a complex system of weights and activation functions.
+#
+# In our case, we represent the function as a linear combination of kernel functions,
+#
+# $$f(x) = \sum_{i=1}^{M} \alpha_{i} k(x_{i}, x)$$
+#
+# where $x_{i}$ are given by our data, and the $\alpha_{i}$ values are our tunable
+# parameters. In effect, we place a Gaussian kernel at each of our data points $x_{i}$,
+# and we want to tune the $\alpha_{i}$ coefficients so that the resulting kernel
+# approximation gives us values that are close to the $y_{i}$ data points we are given.
+
+# %% [markdown]
+# Let's pretend we have some coefficients alpha. Then the kernel approximation is given
+# by,
+
+# %%
+alpha = np.array([[1], [-1], [0.5]])
+
+plt.figure()
+plt.axes()
+plt.plot(X, np.dot(rbf_kernel(X, x_data, sigma=1), alpha))
+plt.show()
+
+# %% [markdown]
+# ### A Practical Example
+#
+# It's easy to see that if we are given enough data, and we appropriately choose the
+# coefficients $\alpha$, then we can closely approximate a function using kernels.
+#
+# As a practical example, let's see how this works for a real function,
+#
+# $$f(x) = x^{3} - 4 x^{2} + 6 x - 24 + \exp(-x)$$
+#
+# where we assume that we are given data $\lbrace (x_{i}, y_{i}) \rbrace_{i=1}^{M}$
+# taken from the function $f$.
+
+# %%
+f = lambda x: x ** 3 - 4 * x ** 2 + 6 * x - 24 + np.exp(-x)
+
+X = np.linspace(-5, 5, 200).reshape(-1, 1)  # Points we want to evaluate at.
+Y = f(X)  # The actual function values. Used for plotting.
+
+# Our observed data.
+sample_size = 5
+x_data = np.linspace(-4, 4, sample_size).reshape(-1, 1)
+y_data = f(x_data)
+
+# %% [markdown]
+# For now, we will tune the parameters manually to approximate the function $f$. We will
+# see shortly, however, that we can compute the coefficients $\alpha$ directly, without
+# having to manually adjust them.
+
+# %%
+alpha = np.array([[-110], [-30], [-20], [-20], [0]])
+
+plt.figure()
+plt.axis()
+plt.plot(X, Y, ":k", label="Actual Function")
+plt.plot(X, np.dot(rbf_kernel(X, x_data, sigma=1), alpha), label="Kernel Approximation")
+plt.legend()
+plt.show()
+
+# %% [markdown]
+# We can see that our approximation is fairly close to the true function, but doesn't
+# approximate it exactly. If we keep adjusting the parameters $\alpha$, we could get an
+# even closer approximation. But tuning the coefficients $\alpha$ by hand can become
+# tedious, especially (as you can imagine) when we have a large amount of data.
+
+# %% [markdown]
+# ### Kernel Learning
+#
+# All learning techniques employ some sort of algorithm to "learn" the tunable
+# parameters in the function representation. For instance, in neural networks we learn
+# the parameters iteratively, by moving them a little bit at a time using a procedure
+# called gradient descent. In order to figure out what direction we need to move the
+# parameters, we define the approximation "error". Then, the gradient of the error tells
+# us which direction to move the parameters. If we take enough small steps, we will
+# (hopefully) converge to the closest possible approximation, or in other words, the
+# solution with the smallest error.
+#
+# This is what we mean by *"learning"*. However, not all learning methods require us to
+# take many small steps to find the "best" values for the parameters or coefficients.
+#
+# __In kernel methods, we don't need to slowly adjust the parameters over time, and we
+# can directly compute the coefficients that give us the best approximation of our
+# function.__
+#
+# This is because in kernel methods, the error is given by a regularized least squares
+# problem, which has a *closed-form solution*.
+
+# %%
+G = rbf_kernel(x_data, sigma=1)
+alpha = np.dot(regularized_inverse(G, regularization_param=1e-5, copy=True), y_data)
+alpha
+
+# %%
+plt.figure()
+plt.axis()
+plt.plot(X, Y, ":k", label="Actual Function")
+plt.plot(X, np.dot(rbf_kernel(X, x_data, sigma=1), alpha), label="Kernel Approximation")
+plt.legend()
+plt.show()
+
+# %% [markdown]
+# We can see that the solution is closer to the true function we came up with when we
+# tuned the coefficient manually, and doesn't require us to manually tune the
+# coefficients $\alpha$. While it may be slightly better or worse in some specific
+# regions, it is a better approximation __overall__.
+#
+# Note that the ``regularization_param`` is another hyperparameter (denoted by $\lambda
+# > 0$) that is introduced by the regularized least-squares problem, and introduces a
+# little extra error in the solution to keep the function smooth. Ideally, we want this
+# value to be as small as possible. Intuitively, with a small amount of data, the
+# solution can "overfit" to the data points, and won't be close to the true function
+# in-between the data points we have. Thus, we use $\lambda$ to prevent overfitting when
+# the dataset is small. However, as the amount of data increases, we can reduce
+# $\lambda$ to be almost zero, and the extra error becomes negligible.
+#
+# Here's the same example with a larger amount of random data.
+
+# %%
+sample_size = 100
+x_data = 8 * np.random.rand(sample_size, 1) - 4  # Random data in (-4, 4)
+y_data = f(x_data)
+
+G = rbf_kernel(x_data, sigma=1)
+alpha = np.dot(regularized_inverse(G, regularization_param=1e-9), y_data)
+
+plt.figure()
+plt.axis()
+plt.scatter(x_data, y_data, s=5, c="red", label="Data")
+plt.plot(X, Y, ":k", label="Actual Function")
+plt.plot(X, np.dot(rbf_kernel(X, x_data, sigma=1), alpha), label="Kernel Approximation")
+plt.legend()
+plt.show()
+
+# %% [markdown]
+# We can see that as the amount of data increases, our approximation gets very close to
+# the true function.
+
+# %% [markdown]
+# ### *What is an RKHS?*
+#
+# A reproducing kernel Hilbert space (RKHS) is effectively a function space defined by
+# the kernel. When we choose a kernel function, we are specifying what the functions in
+# the RKHS look like. For instance, when we choose a Gaussian kernel, the associated
+# RKHS is the set of all functions that can be represented by a linear combination of
+# Gaussians. If we use a polynomial kernel, then the associated RKHS is the set of all
+# polynomials of a certain order.
+#
+# When we dive into the math of reproducing kernel Hilbert spaces, things can become
+# daunting very quickly since it involves a significant amount of functional analysis.
+# This deters many people from using kernel methods or studying them, since the math is
+# considered too "hard". However, this should not be the case, as the concept of an RKHS
+# is superficially very easy to grasp.
+
+# %% [raw] raw_mimetype="text/restructuredtext"
+# Advanced Topics
+# ---------------
+#
+# Here, we've given a basic understanding of kernel methods and how kernels can be used
+# for function approximation. However, we've barely scratched the surface when it comes
+# to the ways in which kernels can be used. For those who are interested in learning
+# more about the techniques used in SOCKS, or simply want to learn more about kernel
+# methods, we've provided some tutorials on some advanced topics. These tutorials are
+# geared toward those who already have a background in machine learning, but may not be
+# as familiar with kernel methods.
+#
+# .. toctree::
+#     :maxdepth: 1
+#
+#     rkhs_tutorial
+#     statistical_learning
+#     nn_comparison
+#
+# Also, be sure to check out the :doc:`examples </examples/index>` page for a showcase
+# of the algorithms in SOCKS.
+#
diff --git a/docs/guide/templates.rst b/docs/guide/templates.rst
index 80020ba..e17505e 100644
--- a/docs/guide/templates.rst
+++ b/docs/guide/templates.rst
@@ -179,7 +179,7 @@ It is recommended to use the following templates when defining new systems
                     v = self.np_random.standard_normal(
                         size=self.observation_space.shape
                     )
-                    return np.array(state, dtype=np.float32) + np.array(v)
+                    return np.array(state, dtype=np.float32) + np.asarray(v)
 
                 def dynamics(self, time, state, action, disturbance):
                     ...
@@ -232,7 +232,7 @@ implements a custom :py:meth:`~gym_socks.envs.policy.BasePolicy.__call__` functi
     The main thing to know when defining custom policies is that functions in SOCKS
     which sample from systems pass all relevant information to the policy for the
     purpose of computing a control action. Since the policy can be either time-varying
-    or time-invariant or open- or closed-lopp, it may require the time and state of the
+    or time-invariant or open- or closed-loop, it may require the time and state of the
     system in order to compute the control action. Thus, custom policies should
     explicitly name these variables as parameters if they are required, and should use
     ``*args, **kwargs`` to capture additional, unneeded arguments.
@@ -253,4 +253,100 @@ algorithms.
             self.action_space = action_space
 
         def __call__(self, *args, **kwargs):
-            return self.action_space.sample()
\ No newline at end of file
+            return self.action_space.sample()
+
+
+Custom Sampling Functions
+=========================
+
+Very simply, a sampling function returns a random sample. In python, this can either be
+a regular function that returns a value or a generator that yields a value. SOCKS
+implements a wrapper function, :py:func:`~gym_socks.sampling.sample.sample_fn`, which
+converts a regular function into an infinite generator.
+
+In order to create a custom sampling function, we need a function that returns or yields
+a single observation, and then we can use the
+:py:func:`~gym_socks.sampling.sample.sample_fn` decorator to convert it into a sample
+function.
+
+Sampling Function Templates
+---------------------------
+
+It is recommended to use the following templates when defining a new sampling function.
+
+.. tab-set::
+
+    .. tab-item:: Basic
+
+        .. code-block:: python
+
+            from gym_socks.sampling import sample_fn
+
+            @sample_fn
+            def custom_sampler():
+
+                # Generate observation here.
+
+                return obs
+
+    .. tab-item:: Generator
+
+        .. code-block:: python
+            :emphasize-lines: 8
+
+            from gym_socks.sampling import sample_fn
+
+            @sample_fn
+            def custom_sampler():
+
+                # Generate observation here.
+
+                yield obs
+
+    .. tab-item:: Partially Observable Sampler
+
+        .. code-block:: python
+
+            from gym_socks.sampling import sample_fn
+            from gym_socks.sampling import space_sampler
+
+            from gym_socks.envs.integrator import NDIntegrator
+            from gym_socks.policies import RandomizedPolicy
+
+            # Partially observable ND integrator system.
+            # Replace this with your own system.
+            class PartiallyObservableNDIntegrator(NDIntegratorEnv):
+                def generate_observation(self, time, state, action):
+                    v = self.np_random.standard_normal(
+                        size=self.observation_space.shape
+                    )
+                    return np.array(state, dtype=np.float32) + np.asarray(v)
+
+            env = PartiallyObservableNDIntegrator(dim=2)
+
+            state_sampler = space_sampler(space=env.state_space)
+            action_sampler = space_sampler(space=env.action_space)
+
+            @sample_fn
+            def custom_sampler():
+                state = next(state_sampler)
+                action = next(action_space)
+
+                env.reset(state)
+                observation, *_ = env.step(action=action)
+                next_state = env.state
+
+                return state, action, next_state, observation
+
+.. note::
+
+    The main reason to use a generator instead of a regular function is if you have
+    internal variables or arguments that are passed to the generator whose value should
+    be preserved or "remembered" between calls.
+
+.. note::
+
+    Behind the scenes, the :py:func:`~gym_socks.sampling.sample.sample_fn` decorator
+    wraps the sample function in a helper class that provides several useful functions
+    that allow us to generate a finite sample and manipulate the generator, for instance
+    to repeat outputs multiple times.
\ No newline at end of file
diff --git a/docs/guide/using_socks.rst b/docs/guide/using_socks.rst
index 7c1e5f6..1547028 100644
--- a/docs/guide/using_socks.rst
+++ b/docs/guide/using_socks.rst
@@ -2,34 +2,161 @@
 Using SOCKS
 ***********
 
+.. toctree::
+    :hidden:
+    :maxdepth: 1
+
+    for_gym_users
+
+
 In practice, we are typically given a dataset to work with, either from an actual system
 or from historical observations of a system. However, as researchers, we typically need
 to generate a sample via simulation in order to develop and test algorithms.
 
+In the following sections, we describe how to organize data that is given to us from a
+pre-existing dataset and how to generate synthetic data by simulating a system.
+
+Organizing Data
+===============
+
+If a dataset is already available, we need to organize and format the data correctly to
+work with the algorithms in SOCKS. In this section, we will describe how to format and
+organize the data into the correct format, and give some general cases.
+
+Data Arrays
+-----------
+
+In a dynamical system, we generally observe transitions from the dynamical system
+equations, meaning our data consists of states :math:`x_{i}`, the applied control
+actions :math:`u_{i}`, and the resulting state after a single time step :math:`y_{i}`.
+
+In python, data is organized in a 2D array with **one data point per row**. We want to
+organize the states into an array ``X``, the control actions into an array ``U``, and
+the resulting states into an array ``Y``. If we have states that are
+:math:`n`-dimensional and :math:`M` sample points, then the arrays ``X`` and ``Y`` will
+have dimensions :math:`M \times n`. Similarly, if the control actions are
+:math:`m`-dimensional, then the array ``U`` will have dimensions :math:`M \times m`.
+
+.. caution::
+
+    Note that this is different from how data is typically ordered in Matlab, where it
+    is easier to order data in *columns* due to the way Matlab orders the elements of
+    matrices. Be careful when importing data from Matlab, since it may need to be
+    transposed in order to fit the correct format.
+
+.. seealso::
+
+    If you have data available in Matlab, and you need to import it into python, check
+    out `scipy.io.loadmat`_.
+
+.. _scipy.io.loadmat:
+    https://docs.scipy.org/doc/scipy/reference/generated/scipy.io.loadmat.html
+
+.. tip::
+
+    If we are dealing with an uncontrolled system, meaning there are no control actions,
+    specify ``U`` as an array of zeros, i.e. ``U = np.zeros(M, 1)``, where ``M`` is the
+    number of sample points.
+
+Mathematically, we typically organize a sample (dataset) as a collection of tuples,
+
+.. math::
+
+    \mathcal{S} = \lbrace (x_{i}, u_{i}, y_{i}) \rbrace_{i=1}^{M}
+
+SOCKS uses the mathematical organization, meaning if we construct a set of 2D arrays, we
+need to convert this to a list of tuples. Luckily, this is a simple operation in python,
+and we can use the built-in ``zip`` function.
+
+Common Cases
+------------
+
+Here we outline some common cases, and give a small code snippet describing how to
+collect and organize the data.
+
+State Transitions
+^^^^^^^^^^^^^^^^^
+
+The most general case is where data is given as state transitions. We have the initial
+states :math:`x_{i}`, the control actions applied to the system :math:`u_{i}`, and the
+next states of the system after a single time step :math:`y_{i}`. The first step in
+organizing the data should be to load all of the data into arrays ``X``, ``U``, and
+``Y``. This part is up to you. Once you have these arrays, we need to convert this to a
+sample :math:`\mathcal{S}`.
+
+If ``X``, ``U``, and ``Y`` are correctly organized, such that each row in the arrays
+corresponds to a single data point, then converting to a sample is straightforward using
+the ``zip`` function.
+
+.. code-block:: python
+
+    S = list(zip(X, U, Y))  # A list of tuples: [(x1, u1, y1), ..., (xn, un, yn)].
+
+Trajectory Data
+^^^^^^^^^^^^^^^
+
+Sometimes, dynamical system data is given as trajectories (meaning a sequence of states
+and control actions indexed by time), and we would like to convert this to a collection
+of state transitions. In this case, we need to break the trajectory down into its
+"increments".
+
+Regardless of how you choose to load the trajectories into python, we typically have a
+2D array of states that are indexed by time. I.e. the trajectory has dimensions :math:`N
+\times n`, where :math:`N` is the number of time steps and :math:`n` is the
+dimensionality of the state space. In addition, we may have a sequence of control
+actions applied at each time step (excluding the last time step), which has dimensions
+:math:`N-1 \times m`, where :math:`m` is the dimensionality of the control or input
+space. In order to split this into increments, we just need to slice the array
+appropriately.
+
+Suppose we are given a state ``trajectory`` as a 2D array. Then in order to separate it
+into the correct format, we can use:
+
+.. code-block:: python
+
+    X = trajectory[:-1]  # Get all rows except the last.
+    Y = trajectory[1:]  # Get all rows except the first.
+
+Then, we can combine these into the correct format for socks using the ``zip`` function
+as described above.
+
+Note that if we have multiple trajectories, we then need to concatenate the states from
+each trajectory, e.g. using `numpy.vstack`_.
+
+.. _numpy.vstack: https://numpy.org/doc/stable/reference/generated/numpy.vstack.html
+
+.. important::
+
+    When splitting a trajectory like this, it is important to make sure that the rows of
+    ``X`` line up with the corresponding control actions in ``U`` and the resulting
+    states in ``Y``. Be careful when concatenating arrays and slicing them to ensure
+    that the state transitions line up correctly.
+
 Simulating Systems
 ==================
 
-SOCKS uses the OpenAI gym framework to simulate systems (called environments in gym).
+In some cases, we may not be dealing with a pre-existing dataset. In this case, we
+typically need to generate an artificial dataset by simulating a system.
+
+SOCKS uses the OpenAI gym framework to simulate systems (called "environments" in gym).
 For example, the following code can be used to simulate a 2D stochastic integrator
 system over 10 time steps.
 
 .. code-block:: python
 
-    import gym
-    import gym_socks
-    env = gym.make("2DIntegratorEnv-v0")
+    from gym_socks.envs.integrator import NDIntegrator
+    env = NDIntegrator(dim=2)  # A 2D stochastic integrator system.
+
+    # Reset the system to a random initial condition.
     env.reset()
-    policy = RandomizedPolicy(action_space=env.action_space)
+
     for t in range(10):
-        action = policy()
+        action = env.action_space.sample()  # Generate a random control action.
         obs, *_ = env.step(time=t, action=action)
 
-The :py:meth:`reset` function resets the system to a random initial condition. We then
-define a randomized policy, which selects a random control action from the action space.
-This is equivalent to using ``env.action_space.sample()``, but we make explicit that the
-policy is a separate object. Then, inside the ``for`` loop, we compute a control action,
-feed it to the dynamics using the :py:meth:`step` function, and obtain an observation of
-the system state ``obs``.
+The :py:meth:`reset` function resets the system to a random initial condition. Then,
+inside the ``for`` loop, we compute a control action, feed it to the dynamics using the
+:py:meth:`step` function, and obtain an observation of the system state ``obs``.
 
 A full list of the environments included in SOCKS is provided in the
 :doc:`/api/gym_socks/envs/index` section of the :doc:`/api/index`.
@@ -77,225 +204,180 @@ A full list of the environments included in SOCKS is provided in the
 
 .. [#ND] The string identifier for these systems generates a 2D system.
 
-
-Sampling Functions
-==================
+Generating a Sample
+-------------------
 
 The main idea of sampling using SOCKS is to define a generator function that yields a
-single observation in the sample :math:`\mathcal{S}`.
+tuple in the sample :math:`\mathcal{S}`,
+
+.. math::
+
+    \mathcal{S} = \lbrace (x_{i}, u_{i}, y_{i}) \rbrace_{i=1}^{M}
 
 .. note::
 
-    A generator in Python can be thought of as a function that remembers its state and
+    A generator in python can be thought of as a function that remembers its state and
     may be called multiple times to return a sequence of values. In our case, we use it
     to "generate" a sequence of observations.
 
-.. hint::
+.. tip::
 
     An infinite generator can be used directly, since it can be used to generate a
     sample of any length. However, not all generator functions are easily defined as
     infinite generators. A finite generator (or even a non-generator function that
     returns an observation) can still be used, and SOCKS provides a decorator called
-    :py:func:`sample_generator` that wraps a sample function, effectively turning it
+    :py:func:`sample_fn` that wraps a sample function, effectively turning it
     into an infinite generator.
 
 The important points to remember:
 
 * The states :math:`x_{i}` are drawn from a distribution on :math:`\mathcal{X}`.
 * The actions :math:`u_{i}` are drawn from a distribution on :math:`\mathcal{U}`.
-* The resulting states :math:`x_{i}{}^{\prime}` are drawn from a stochastic kernel.
+* The resulting states :math:`y_{i}` are drawn from a stochastic kernel.
 
-Thus, we define a sample generator for states and actions, and another to generate the
-tuple :math:`(x_{i}, u_{i}, x_{i}{}^{\prime})`. Very simply, a sample generator for
-actions could be implemented as :py:obj:`action_space.sample()`, which is implemented in
-:py:obj:`gym`.
+In order to generate a sample, we need a way to generate random states, compute control
+actions (which can either be chosen randomly or chosen via a policy), and compute the
+resulting states. SOCKS implements several functions to make this process easier. In a
+nutshell, we need to define a sample generator for states and actions, and another to
+generate the tuple :math:`(x_{i}, u_{i}, y_{i})`.
 
-Sampling Function Examples
---------------------------
+In order to sample from a space, we can use the
+:py:func:`~gym_socks.sampling.sample.space_sampler` function, which randomly generates
+samples from a :py:obj:`gym.Space`.
 
-.. tab-set::
-
-    .. tab-item:: Basic
-
-        .. code-block:: python
-
-            state_sampler = random_sampler(sample_space=env.state_space)
-            action_sampler = random_sampler(sample_space=env.action_space)
-
-            @sample_generator
-            def custom_sampler():
-                state = next(state_sampler)
-                action = next(action_sampler)
-
-                env.state = state
-                next_state, *_ = env.step(action=action)
-
-                yield (state, action, next_state)
-
-    .. tab-item:: Policy
-
-        .. code-block:: python
-
-            state_sampler = random_sampler(sample_space=env.state_space)
-            policy = RandomizedPolicy(action_space=env.action_space)
-
-            @sample_generator
-            def custom_sampler():
-                state = next(state_sampler)
-                action = policy()
+.. code-block:: python
 
-                env.state = state
-                next_state, *_ = env.step(action=action)
+    from gym.spaces import Box
+    from gym_socks.sampling import space_sampler
 
-                yield (state, action, next_state)
+    sample_space = Box(low=-1, high=1, shape(2,), dtype=float)
 
+    # Generate a single observation from the space.
+    x = next(space_sampler)
 
-Making Experiments
-==================
+    # Generate a sample of 100 observations.
+    X = space_sampler(space=sample_space).sample(100)
 
-We strongly encourage using the `sacred <https://github.com/IDSIA/sacred>`_ experimental
-framework for new experiments. This promotes repeatability and easy configuration that
-is useful for rapid development, iteration, and modification.
+.. seealso::
 
-Minimal Experiment
-------------------
+    SOCKS also provides a sampling function which generates points from a pre-specified
+    grid, :py:func:`~gym_socks.sampling.sample.grid_sampler`. This can be useful in
+    certain cases to obtain a more uniform result.
 
-The following template is a (mostly) minimal example of how to organize a new
-experiment, and allows for simple modification and extension.
+We can use the same procedure to randomly sample from the state and action spaces of an
+environment.
 
 .. code-block:: python
 
-    import gym
-    import gym_socks
-    import logging
-
-    import numpy as np
-
-    from sacred import Experiment
-
-    from examples._computation_timer import ComputationTimer
-
-    from examples.ingredients.system_ingredient import system_ingredient
-    from examples.ingredients.system_ingredient import set_system_seed
-    from examples.ingredients.system_ingredient import make_system
+    from gym_socks.envs.integrator import NDIntegratorEnv
+    from gym_socks.sampling import space_sampler
 
+    env = NDIntegratorEnv(2)
+    state_sampler = space_sampler(env.state_space)
+    action_sampler = space_sampler(env.action_space)
 
-    @system_ingredient.config
-    def system_config():
-        system_id = "2DIntegratorEnv-v0"
-        sampling_time = 0.25
-
-
-    ex = Experiment(ingredients=[system_ingredient])
+Then, SOCKS implements a function
+:py:func:`~gym_socks.sampling.sample.transition_sampler`, which can be used to generate
+the tuples in :math:`\mathcal{S}`. Continuing from the last code block,
 
+.. code-block:: python
 
-    @ex.config
-    def config():
-        time_horizon = 5
-        results_filename = "results/data.npy"
+    from gym_socks.sampling import transition_sampler
 
+    # Generate a sample of 100 observations.
+    S = transition_sampler(env, state_sampler, action_sampler).sample(100)
 
-    @ex.main
-    def main(seed, time_horizon, results_filename, _log):
+.. seealso::
 
-        env = make_system()
+    There is also a sampling function,
+    :py:func:`~gym_socks.sampling.sample.trajectory_sampler`, which generates samples
+    consisting of initial conditions, control sequences, and resulting trajectories.
+    This can be useful depending on the algorithm you are using.
 
-        # Set the random seed.
-        set_system_seed(seed, env)
+Custom Sampling Functions
+-------------------------
 
-        with ComputationTimer():
+Sometimes, the built-in sampling functions in SOCKS may be insufficient for your needs.
+If you need to manipulate the way the samples are generated, for instance to change the
+sampling distribution, then you may need to define your own sampling function. This is
+achievable in SOCKS using the :py:func:`~gym_socks.sampling.sample.sample_fn` decorator.
 
-            # Run algorithm here.
+.. tab-set::
 
-        # Save the result to NPY file.
-        _log.debug(f"Saving the results to file {results_filename}.")
-        with open(results_filename, "wb") as f:
-            np.save(f, result)
+    .. tab-item:: Basic
 
+        .. code-block:: python
 
-    if __name__ == "__main__":
-        ex.run_commandline()
+            from gym_socks.sampling import sample_fn
+            from gym_socks.sampling import space_sampler
 
-The "magic" of sacred comes from the configuration functions, and how the variables
-defined within them can be updated and injected into the ``main`` function. However, the
-main ideas are to define parameters and configuration variables within the ``config``
-function, and to add them as arguments to the ``main`` function. Then, at runtime,
-sacred will pass either the default value specified in the config, or use the updated
-values passed by the user at the command line. In fact, this works for any "captured"
-function, which means that only the relevant configuration variables need to be passed
-to the functions that use them. This allows for more modularized experiments and also
-allows for more convenient interaction by the user.
+            from gym_socks.envs.integrator import NDIntegrator
 
-Adding Plotting
----------------
+            env = NDIntegrator(dim=2)  # A 2D stochastic integrator system.
 
-Plotting results is typically desired for many experiments, and the sacred framework can
-also be used to create configurable plotting commands.
+            state_sampler = space_sampler(space=env.state_space)
+            action_sampler = space_sampler(space=env.action_space)
 
-The following code can be added to the experiment, and then a plot can be generated
-using the command:
+            @sample_fn
+            def custom_sampler():
+                state = next(state_sampler)
+                action = next(action_sampler)
 
-.. code-block:: shell
+                env.reset(state)
+                next_state, *_ = env.step(action=action)
 
-    python -m <experiment> plot_results
+                return state, action, next_state
 
-.. code-block:: python
+            # Generate a sample of 100 observations.
+            S = custom_sampler().sample(size=100)
 
-    from examples.ingredients.plotting_ingredient import plotting_ingredient
-    from examples.ingredients.plotting_ingredient import update_rc_params
+    .. tab-item:: Policy
 
-    @plotting_ingredient.config_hook
-    def _plot_config(config, command_name, logger):
-        if command_name in {"main", "plot_results"}:
-            return {
-                "axes": {
-                    "xlabel": r"$x_1$",
-                    "ylabel": r"$x_2$",
-                }
-            }
+        .. code-block:: python
+            :emphasize-lines: 5, 10, 15
 
+            from gym_socks.sampling import sample_fn
+            from gym_socks.sampling import space_sampler
 
-    @ex.command(unobserved=True)
-    def plot_results(plot_cfg, results_filename):
+            from gym_socks.envs.integrator import NDIntegrator
+            from gym_socks.policies import RandomizedPolicy
 
-        logging.getLogger("matplotlib").setLevel(logging.WARNING)
-        logging.getLogger("PIL").setLevel(logging.WARNING)
+            env = NDIntegrator(dim=2)  # A 2D stochastic integrator system.
 
-        # Dynamically load for speed.
-        import matplotlib
+            state_sampler = space_sampler(space=env.state_space)
+            policy = RandomizedPolicy(action_space=env.action_space)
 
-        matplotlib.use("Agg")
-        update_rc_params(matplotlib, plot_cfg["rc_params"])
+            @sample_fn
+            def custom_sampler():
+                state = next(state_sampler)
+                action = policy(state=state)
 
-        import matplotlib.pyplot as plt
+                env.reset(state)
+                next_state, *_ = env.step(action=action)
 
-        # Load the result from NPY file.
-        with open(results_filename, "rb") as f:
-            result = np.load(f)
+                return state, action, next_state
 
-        fig = plt.figure()
-        ax = plt.axes(**plot_cfg["axes"])
+            # Generate a sample of 100 observations.
+            S = custom_sampler().sample(size=100)
 
-        # Plotting code here.
+Check out the :doc:`/guide/templates` page for some templates and a description of how
+you can create your own sampling functions.
 
-        plt.savefig(plot_cfg["plot_filename"])
+Algorithms
+==========
 
-The utility of this approach is when the format of the results needs to be changed to
-fit a particular publication, but re-running the experiment can be time-intensive. By
-saving the main algorithm results to a file, and then loading them separately in the
-plotting function, we save the time of having to re-compute the algorithm for plot
-manipulation. The :py:mod:`~examples.ingredients.plotting_ingredient` module implements
-a small ingredient that can be used to modify the ``rc_params`` of ``matplotlib``, or
-add configuration options that can be used to modify the plot appearance. For example,
+The algorithms in SOCKS are mostly designed to follow the format of `scikit-learn`_,
+meaning they typically have some kind of "fit/predict" function to train the algorithm
+and then evaluate the solution.
 
-.. code-block:: shell
+.. _scikit-learn: https://scikit-learn.org/stable/index.html
 
-    python -m <experiment> plot_results with <updates>
+Because the algorithms handle a number of different and distinct problems, we do not
+describe them here. Instead, check out the :doc:`/examples/index` page to get an idea of
+how they work and how to use them.
 
-.. hint::
+If you have questions or get stuck, feel free to reach out by posting an `issue
+<https://github.com/ajthor/socks/issues>`_ on GitHub or by starting a `discussion
+<https://github.com/ajthor/socks/discussions>`_.
 
-    The plotting config is typically hidden when using the ``print_config`` command
-    using the above code. This is due to the ``if command_name in {"main",
-    "plot_results"}`` line above. To show the config, either use the ``-p`` flag when
-    calling the ``plot_results`` command, or remove the ``if`` statement in the
-    ``_plot_config`` function above.
\ No newline at end of file
+We'd love to hear from you!
\ No newline at end of file
diff --git a/docs/index.rst b/docs/index.rst
index 2294678..75992b8 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -31,6 +31,7 @@ environments.
     :caption: User Guide
 
     guide/fundamentals
+    guide/kernel_tutorial/tutorial
     guide/using_socks
     guide/templates
     contributing/index
@@ -57,6 +58,9 @@ Check out the user guide to use SOCKS in your own projects.
     Describes the basic concepts of SOCKS, and gives a quick overview of data-driven
     control.
 
+:doc:`guide/kernel_tutorial/tutorial`
+    A tutorial on kernel methods to get you started.
+
 :doc:`guide/using_socks`
     Information on how to simulate and generate samples from systems, as well as basic
     information about how to use the algorithms in SOCKS.
@@ -143,9 +147,13 @@ In order to cite the toolbox, use the following bibtex entry:
       title     = {{SOCKS}: A Kernel-Based Stochastic Optimal Control and Reachability Toolbox},
       authors   = {Thorpe, Adam J. and Oishi, Meeko M. K.},
       year      = {2022},
-      booktitle = {Proceedings of the 25th ACM International Conference on Hybrid Systems: Computation and Control (submitted)}
+      booktitle = {Proceedings of the 25th ACM International Conference on Hybrid Systems: Computation and Control}
     }
 
+Or from `ACM <https://doi.org/10.1145/3501710.3519525>`_.
+
+Also check out the live presentation of SOCKS on `YouTube
+<https://youtu.be/EfIPzpHy-YQ>`_.
 
 Indices and Tables
 ==================
diff --git a/examples/control/satellite_rendezvous.py b/examples/control/satellite_rendezvous.py
index a411020..77c482b 100644
--- a/examples/control/satellite_rendezvous.py
+++ b/examples/control/satellite_rendezvous.py
@@ -25,13 +25,12 @@
 
 from gym_socks.algorithms.control.kernel_control_fwd import KernelControlFwd
 
-from sklearn.metrics.pairwise import rbf_kernel
+from gym_socks.kernel.metrics import rbf_kernel
 
-from gym_socks.sampling import sample
-from gym_socks.sampling import default_sampler
-from gym_socks.sampling import random_sampler
+from gym_socks.sampling import space_sampler
+from gym_socks.sampling import transition_sampler
 
-from gym_socks.utils.grid import make_grid_from_space
+from gym_socks.utils.grid import boxgrid
 
 
 # %% [markdown]
@@ -55,7 +54,6 @@
 # %%
 env = make(system_id)
 env.seed(seed)
-
 sample_size = 2500
 
 state_sample_space = gym.spaces.Box(
@@ -65,7 +63,6 @@
     dtype=np.float32,
     seed=seed,
 )
-state_sampler = random_sampler(sample_space=state_sample_space)
 
 action_sample_space = gym.spaces.Box(
     low=-0.05,
@@ -74,18 +71,12 @@
     dtype=np.float32,
     seed=seed,
 )
-action_sampler = random_sampler(sample_space=action_sample_space)
-
-S = sample(
-    sampler=default_sampler(
-        state_sampler=state_sampler,
-        action_sampler=action_sampler,
-        env=env,
-    ),
-    sample_size=sample_size,
-)
 
-A = make_grid_from_space(sample_space=action_sample_space, resolution=[20, 20])
+state_sampler = space_sampler(state_sample_space)
+action_sampler = space_sampler(action_sample_space)
+S = transition_sampler(env, state_sampler, action_sampler).sample(size=sample_size)
+
+A = boxgrid(space=action_sample_space, resolution=[20, 20])
 
 
 # %% [markdown]
@@ -186,11 +177,11 @@ def _constraint_fn(time: int = 0, state: np.ndarray = None) -> float:
 # Simulate the controlled system.
 env.reset()
 initial_condition = [-0.8, -0.8, 0, 0]
-env.state = initial_condition
+env.reset(initial_condition)
 trajectory = [initial_condition]
 
 for t in range(time_horizon):
-    action = policy(time=t, state=[env.state])
+    action = policy(time=t, state=env.state)
     state, *_ = env.step(time=t, action=action)
 
     trajectory.append(list(state))
diff --git a/examples/control/tracking.py b/examples/control/tracking.py
index 50a0d4c..6a945ba 100644
--- a/examples/control/tracking.py
+++ b/examples/control/tracking.py
@@ -25,14 +25,13 @@
 from gym_socks.algorithms.control.kernel_control_bwd import KernelControlBwd
 
 from functools import partial
-from sklearn.metrics.pairwise import rbf_kernel
 
-from gym_socks.sampling import sample
-from gym_socks.sampling import default_sampler
-from gym_socks.sampling import random_sampler
-from gym_socks.sampling import grid_sampler
+from gym_socks.kernel.metrics import rbf_kernel
 
-from gym_socks.utils.grid import make_grid_from_ranges
+from gym_socks.sampling import space_sampler
+from gym_socks.sampling import transition_sampler
+
+from gym_socks.utils.grid import cartesian
 
 # %% [markdown]
 # Configuration variables.
@@ -58,34 +57,29 @@
 env = make(system_id)
 env.sampling_time = 0.1
 env.seed(seed)
-env.action_space = gym.spaces.Box(
-    low=np.array([0.1, -10.1], dtype=np.float32),
-    high=np.array([1.1, 10.1], dtype=np.float32),
-    shape=(2,),
-    dtype=np.float32,
-    seed=seed,
-)
-
 sample_size = 1500
 
-sample_space = gym.spaces.Box(
+state_sample_space = gym.spaces.Box(
     low=np.array([-1.2, -1.2, -2 * np.pi], dtype=np.float32),
     high=np.array([1.2, 1.2, 2 * np.pi], dtype=np.float32),
     shape=(3,),
     dtype=np.float32,
     seed=seed,
 )
-state_sampler = random_sampler(sample_space=sample_space)
-action_sampler = random_sampler(sample_space=env.action_space)
-
-S = sample(
-    sampler=default_sampler(
-        state_sampler=state_sampler, action_sampler=action_sampler, env=env
-    ),
-    sample_size=sample_size,
+
+action_sample_space = gym.spaces.Box(
+    low=np.array([0.1, -10.1], dtype=np.float32),
+    high=np.array([1.1, 10.1], dtype=np.float32),
+    shape=(2,),
+    dtype=np.float32,
+    seed=seed,
 )
 
-A = make_grid_from_ranges([np.linspace(0.1, 1.1, 10), np.linspace(-10.1, 10.1, 21)])
+state_sampler = space_sampler(state_sample_space)
+action_sampler = space_sampler(action_sample_space)
+S = transition_sampler(env, state_sampler, action_sampler).sample(size=sample_size)
+
+A = cartesian(np.linspace(0.1, 1.1, 10), np.linspace(-10.1, 10.1, 21))
 
 # %% [markdown]
 # We define the cost as the norm distance to the target at each time step.
@@ -136,7 +130,7 @@ def _tracking_cost(time: int = 0, state: np.ndarray = None) -> float:
 policy = KernelControlFwd(
     time_horizon=time_horizon,
     cost_fn=_tracking_cost,
-    kernel_fn=partial(rbf_kernel, gamma=1 / (2 * (sigma ** 2))),
+    kernel_fn=partial(rbf_kernel, sigma=sigma),
     regularization_param=regularization_param,
     verbose=False,
 )
@@ -146,11 +140,11 @@ def _tracking_cost(time: int = 0, state: np.ndarray = None) -> float:
 # Simulate the controlled system.
 env.reset()
 initial_condition = [-0.8, 0, 0]
-env.state = initial_condition
+env.reset(initial_condition)
 trajectory = [initial_condition]
 
 for t in range(time_horizon):
-    action = policy(time=t, state=[env.state])
+    action = policy(time=t, state=env.state)
     state, *_ = env.step(time=t, action=action)
 
     trajectory.append(list(state))
diff --git a/examples/experimental/invariant_transform.py b/examples/experimental/invariant_transform.py
index e8d7301..a799fc3 100644
--- a/examples/experimental/invariant_transform.py
+++ b/examples/experimental/invariant_transform.py
@@ -49,11 +49,11 @@
 
 from functools import partial
 from itertools import islice
-from sklearn.metrics.pairwise import rbf_kernel
 
-from gym_socks.sampling import sample
-from gym_socks.sampling import random_sampler
-from gym_socks.sampling import sample_generator
+from gym_socks.kernel.metrics import rbf_kernel
+
+from gym_socks.sampling import sample_fn
+from gym_socks.sampling import space_sampler
 from gym_socks.sampling.transform import transpose_sample
 
 import matplotlib
@@ -72,15 +72,14 @@ def generate_disturbance(self, time, state, action):
 
 deterministic_env = DeterministicNonholonomicVehicleEnv()
 
+# Generate a collection of random actions.
 action_sample_space = gym.spaces.Box(
     low=np.array([0.1, -10.1], dtype=np.float32),
     high=np.array([1.1, 10.1], dtype=np.float32),
     shape=deterministic_env.action_space.shape,
     dtype=deterministic_env.action_space.dtype,
 )
-action_sampler = random_sampler(action_sample_space)
-# Generate a collection of random actions.
-random_actions = list(islice(action_sampler, 100))
+random_actions = space_sampler(action_sample_space).sample(size=100)
 
 
 # %% [markdown]
@@ -97,14 +96,14 @@ def generate_disturbance(self, time, state, action):
 resulting_states1 = []
 initial_condition1 = [0, 0, 0]
 for action in random_actions:
-    deterministic_env.state = initial_condition1
+    deterministic_env.reset(initial_condition1)
     deterministic_env.step(action=action)
     resulting_states1.append(deterministic_env.state)
 
 resulting_states2 = []
 initial_condition2 = [0.2, 0, 0]
 for action in random_actions:
-    deterministic_env.state = initial_condition2
+    deterministic_env.reset(initial_condition2)
     deterministic_env.step(action=action)
     resulting_states2.append(deterministic_env.state)
 
@@ -133,14 +132,14 @@ def generate_disturbance(self, time, state, action):
 resulting_states1 = []
 initial_condition1 = [0, 0, 0]
 for action in random_actions:
-    deterministic_env.state = initial_condition1
+    deterministic_env.reset(initial_condition1)
     deterministic_env.step(action=action)
     resulting_states1.append(deterministic_env.state)
 
 resulting_states2 = []
 initial_condition2 = [0, 0, np.pi / 2]
 for action in random_actions:
-    deterministic_env.state = initial_condition2
+    deterministic_env.reset(initial_condition2)
     deterministic_env.step(action=action)
     resulting_states2.append(deterministic_env.state)
 
@@ -167,7 +166,7 @@ def generate_disturbance(self, time, state, action):
 # %%
 sigma = 1.5  # Kernel bandwidth parameter.
 regularization_param = 1e-5  # Regularization parameter.
-kernel_fn = partial(rbf_kernel, gamma=1 / (2 * sigma ** 2))
+kernel_fn = partial(rbf_kernel, sigma=sigma)
 
 time_horizon = 20
 
@@ -180,7 +179,7 @@ def generate_disturbance(self, time, state, action):
 # information while we generate it from random initial conditions.
 #
 # _NOTE_: We use significantly less sample information to obtain a similar result to the
-# `tracking` example. Here, since we can translate and rotate the sample, we use 50
+# `tracking` example. Here, since we can translate and rotate the sample, we use 100
 # sample points, wehre in the other example we use 1,500 by default to obtain similar
 # behavior.
 
@@ -197,28 +196,23 @@ def generate_disturbance(self, time, state, action):
     dtype=env.action_space.dtype,
     seed=seed,
 )
-action_sampler = random_sampler(action_sample_space)
+action_sampler = space_sampler(action_sample_space)
 
 
-@sample_generator
+@sample_fn
 def custom_sampler():
-    state = [0, 0, 0]
+    state = np.array([0.0, 0.0, 0.0])
     action = next(action_sampler)
 
-    env.state = state
+    env.reset(state)
     env.step(action=action)
     next_state = env.state
 
-    yield (state, action, next_state)
-
-
-@sample_generator
-def custom_action_sampler():
-    yield next(action_sampler)
+    yield state, action, next_state
 
 
-S = sample(sampler=custom_sampler, sample_size=sample_size)
-A = sample(sampler=custom_action_sampler, sample_size=sample_size)
+S = custom_sampler().sample(size=sample_size)
+A = action_sampler.sample(size=sample_size)
 
 X, U, Y = transpose_sample(S)
 
@@ -311,7 +305,7 @@ def _tracking_cost(time: int = 0, state: np.ndarray = None) -> float:
 # %%
 env.reset()
 initial_condition = [-0.8, 0, 0]
-env.state = initial_condition
+env.reset(initial_condition)
 trajectory = [initial_condition]
 
 for t in range(time_horizon):
diff --git a/examples/identification/linear_id.py b/examples/identification/linear_id.py
index 1674adf..4f5e7a2 100644
--- a/examples/identification/linear_id.py
+++ b/examples/identification/linear_id.py
@@ -26,9 +26,9 @@
 from gym_socks.algorithms.identification.kernel_linear_id import kernel_linear_id
 
 from gym_socks.policies import ConstantPolicy
-from gym_socks.sampling import sample
-from gym_socks.sampling import sample_generator
-from gym_socks.sampling import random_sampler
+
+from gym_socks.sampling import space_sampler
+from gym_socks.sampling import transition_sampler
 
 # %% [markdown]
 # Configuration variables.
@@ -50,23 +50,9 @@
 
 sample_size = 100
 
-state_sampler = random_sampler(sample_space=env.state_space)
-action_sampler = random_sampler(sample_space=env.action_space)
-
-
-@sample_generator
-def sampler():
-    state = next(state_sampler)
-    action = next(action_sampler)
-
-    env.state = state
-    next_state, *_ = env.step(action=action)
-
-    yield (state, action, next_state)
-
-
-S = sample(sampler=sampler, sample_size=sample_size)
-
+state_sampler = space_sampler(env.state_space)
+action_sampler = space_sampler(env.action_space)
+S = transition_sampler(env, state_sampler, action_sampler).sample(size=sample_size)
 
 # %% [markdown]
 # ## Algorithm
@@ -93,10 +79,10 @@ def sampler():
 initial_condition = [-0.75, -0.75, 0, 0]
 
 # Simulate the system using the actual dynamics.
-env.state = initial_condition
+env.reset(initial_condition)
 actual_trajectory = [env.state]
 for t in range(time_horizon):
-    action = np.array(policy(time=t, state=[env.state]), dtype=np.float32)
+    action = policy(time=t, state=env.state)
     obs, *_ = env.step(time=t, action=action)
     next_state = env.state
 
@@ -105,7 +91,7 @@ def sampler():
 # Simulate the system using the approximated dynamics.
 estimated_trajectory = [initial_condition]
 for t in range(time_horizon):
-    action = policy(time=t, state=[env.state])
+    action = policy(time=t, state=env.state)
     state = alg.predict(T=estimated_trajectory[t], U=action)
 
     estimated_trajectory.append(state)
diff --git a/examples/kernel/conditional_embedding.py b/examples/kernel/conditional_embedding.py
index bbdc70a..122ecc0 100644
--- a/examples/kernel/conditional_embedding.py
+++ b/examples/kernel/conditional_embedding.py
@@ -18,7 +18,7 @@
 import numpy as np
 from functools import partial
 from gym_socks.algorithms.kernel import ConditionalEmbedding
-from sklearn.metrics.pairwise import rbf_kernel
+from gym_socks.kernel.metrics import rbf_kernel
 
 from time import perf_counter
 
@@ -45,8 +45,7 @@
 
 # %%
 sigma = 1
-gamma = 1 / (2 * sigma ** 2)
-kernel_fn = partial(rbf_kernel, gamma=gamma)
+kernel_fn = partial(rbf_kernel, sigma=sigma)
 regularization_param = 1 / (sample_size ** 2)
 
 # %% [markdown]
diff --git a/examples/kernel/derivative_approximation.py b/examples/kernel/derivative_approximation.py
index 51c031e..1fb90e9 100644
--- a/examples/kernel/derivative_approximation.py
+++ b/examples/kernel/derivative_approximation.py
@@ -25,8 +25,7 @@
 from gym_socks.algorithms.kernel import ConditionalEmbedding
 from gym_socks.kernel.metrics import rbf_kernel_derivative
 
-from sklearn.metrics.pairwise import rbf_kernel
-from sklearn.metrics.pairwise import euclidean_distances
+from gym_socks.kernel.metrics import rbf_kernel
 
 from time import perf_counter
 
@@ -51,12 +50,11 @@
 
 # %%
 sigma = 0.1
-gamma = 1 / (2 * sigma ** 2)
-kernel_fn = partial(rbf_kernel, gamma=gamma)
+kernel_fn = partial(rbf_kernel, sigma=sigma)
 regularization_param = 1 / (sample_size ** 2)
 
 # %% [markdown]
-# ## Compute the Appromation
+# ## Compute the Approximation
 
 #  %%
 start = perf_counter()
@@ -65,8 +63,8 @@
 K = kernel_fn(X_train, X_test)
 alg.fit(G)
 
-C = rbf_kernel_derivative(X_train, sigma=sigma, distance_fn=euclidean_distances)
-D = rbf_kernel_derivative(X_train, X_test, sigma=sigma, distance_fn=euclidean_distances)
+C = rbf_kernel_derivative(X_train, sigma=sigma)
+D = rbf_kernel_derivative(X_train, X_test, sigma=sigma)
 
 y_pred = alg.predict(y_train, K)
 y_pred_d1 = -alg.predict(y_train, K * D)
diff --git a/examples/kernel/maximum_mean_discrepancy.py b/examples/kernel/maximum_mean_discrepancy.py
index 76ac5a8..b0d72e3 100644
--- a/examples/kernel/maximum_mean_discrepancy.py
+++ b/examples/kernel/maximum_mean_discrepancy.py
@@ -19,7 +19,7 @@
 from gym_socks.kernel.probability import maximum_mean_discrepancy
 from gym_socks.kernel.probability import witness_function
 
-from sklearn.metrics.pairwise import rbf_kernel
+from gym_socks.kernel.metrics import rbf_kernel
 from functools import partial
 
 # %% [markdown]
@@ -33,7 +33,7 @@
 
 # %%
 sigma = 0.25
-kernel_fn = partial(rbf_kernel, gamma=1 / (2 * (sigma ** 2)))
+kernel_fn = partial(rbf_kernel, sigma=sigma)
 
 m = 5000  # sample size for P
 n = 5000  # sample size for Q
diff --git a/examples/kernel/nystrom_approximation.py b/examples/kernel/nystrom_approximation.py
index e8c541d..5858927 100644
--- a/examples/kernel/nystrom_approximation.py
+++ b/examples/kernel/nystrom_approximation.py
@@ -27,7 +27,7 @@
 from functools import partial
 from gym_socks.algorithms.kernel import ConditionalEmbedding
 
-from sklearn.metrics.pairwise import rbf_kernel
+from gym_socks.kernel.metrics import rbf_kernel
 from sklearn.kernel_approximation import Nystroem
 
 from time import perf_counter
@@ -57,8 +57,8 @@
 
 # %%
 sigma = 1
-gamma = 1 / (2 * sigma ** 2)
-kernel_fn = partial(rbf_kernel, gamma=gamma)
+gamma = 1 / (2 * (sigma ** 2))
+kernel_fn = partial(rbf_kernel, sigma=sigma)
 regularization_param = 1 / (sample_size ** 2)
 num_features = 100  # Number of features.
 
diff --git a/examples/kernel/random_fourier_features.py b/examples/kernel/random_fourier_features.py
index 8b45b0b..43d0662 100644
--- a/examples/kernel/random_fourier_features.py
+++ b/examples/kernel/random_fourier_features.py
@@ -20,7 +20,7 @@
 from functools import partial
 from gym_socks.algorithms.kernel import ConditionalEmbedding
 
-from sklearn.metrics.pairwise import rbf_kernel
+from gym_socks.kernel.metrics import rbf_kernel
 from sklearn.kernel_approximation import RBFSampler
 
 from time import perf_counter
@@ -50,8 +50,8 @@
 
 # %%
 sigma = 1
-gamma = 1 / (2 * sigma ** 2)
-kernel_fn = partial(rbf_kernel, gamma=gamma)
+gamma = 1 / (2 * (sigma ** 2))
+kernel_fn = partial(rbf_kernel, sigma=sigma)
 regularization_param = 1 / (sample_size ** 2)
 num_features = 100  # Number of Fourier features.
 
diff --git a/examples/reach/forward_reach.py b/examples/reach/forward_reach.py
index 1d9e3a0..935573d 100644
--- a/examples/reach/forward_reach.py
+++ b/examples/reach/forward_reach.py
@@ -16,23 +16,17 @@
 """
 
 # %%
-import gym_socks
-
 import numpy as np
+from functools import partial
 
-from gym.envs.registration import make
+import matplotlib
+import matplotlib.pyplot as plt
 
 from gym_socks.algorithms.reach.separating_kernel import SeparatingKernelClassifier
 
-from functools import partial
-from sklearn.metrics.pairwise import euclidean_distances
-
-from gym_socks.utils.grid import make_grid_from_ranges
-
-from gym_socks.policies import ConstantPolicy
-from gym_socks.sampling import sample
-from gym_socks.sampling import sample_generator
-from gym_socks.sampling import random_sampler
+from gym_socks.kernel.metrics import abel_kernel
+from gym_socks.sampling import sample_fn
+from gym_socks.utils.grid import cartesian
 
 # %% [markdown]
 # ## Generate The Sample
@@ -41,13 +35,7 @@
 # uniformly within a toroidal region centered around the origin.
 
 # %%
-sigma = 0.1
-sample_size = 1000
-
-regularization_param = 1 / sample_size
-
-
-@sample_generator
+@sample_fn
 def sampler() -> tuple:
     """Sample generator.
 
@@ -67,7 +55,8 @@ def sampler() -> tuple:
 
 
 # Sample the distribution.
-S = sample(sampler=sampler, sample_size=sample_size)
+S = sampler().sample(size=1000)
+S = np.array(S)
 
 # %% [markdown]
 # ## Algorithm
@@ -77,17 +66,12 @@ def sampler() -> tuple:
 
 # %%
 # Construct the algorithm.
-alg = SeparatingKernelClassifier(
-    kernel_fn=partial(
-        gym_socks.kernel.metrics.abel_kernel,
-        sigma=sigma,
-        distance_fn=euclidean_distances,
-    ),
-    regularization_param=regularization_param,
-)
-
-# Generate test points.
-T = make_grid_from_ranges([np.linspace(-1, 1, 50), np.linspace(-1, 1, 50)])
+alg = SeparatingKernelClassifier(kernel_fn=partial(abel_kernel, sigma=0.1))
+
+# Generate evaluation (test) points.
+x1 = np.linspace(-1, 1, 50)
+x2 = np.linspace(-1, 1, 50)
+T = cartesian(x1, x2)
 
 # Train the classifier and classify the points.
 labels = alg.fit(S).predict(T)
@@ -101,18 +85,16 @@ def sampler() -> tuple:
 # countour plot that relies upon a convex hull.
 
 # %%
-import matplotlib
-import matplotlib.pyplot as plt
+# Reshape data.
+XX, YY = np.meshgrid(x1, x2, indexing="ij")
+Z = labels.reshape(XX.shape)
 
+# Plot original data and classified points.
 fig = plt.figure()
 ax = plt.axes()
 
-points_in = T[labels == True]
-
-size = (300.0 / fig.dpi) ** 2
-plt.scatter(points_in[:, 0], points_in[:, 1], color="C0", marker=",", s=size)
-S = np.array(S)
-plt.scatter(S[:, 0], S[:, 1], color="r", marker=".", s=1)
+plt.pcolor(XX, YY, Z, cmap="Blues", shading="auto")
+plt.scatter(S[:, 0], S[:, 1], color="r", marker=".", s=5)
 
 # Plot support region.
 plt.gca().add_patch(plt.Circle((0, 0), 0.5, fc="none", ec="blue"))
diff --git a/examples/reach/stoch_reach.py b/examples/reach/stoch_reach.py
index 4fe9969..afb1e27 100644
--- a/examples/reach/stoch_reach.py
+++ b/examples/reach/stoch_reach.py
@@ -22,17 +22,15 @@
 from functools import partial
 
 from gym.envs.registration import make
-from sklearn.metrics.pairwise import rbf_kernel
+from gym_socks.kernel.metrics import rbf_kernel
 
 from gym_socks.algorithms.reach.kernel_sr import kernel_sr
 
-from gym_socks.sampling import sample
-from gym_socks.sampling import default_sampler
+from gym_socks.sampling import transition_sampler
 from gym_socks.sampling import grid_sampler
-from gym_socks.sampling import repeat
 
-from gym_socks.utils.grid import make_grid_from_space
-from gym_socks.utils.grid import make_grid_from_ranges
+from gym_socks.utils.grid import boxgrid
+from gym_socks.utils.grid import cartesian
 
 # %% [markdown]
 # Configuration variables.
@@ -56,22 +54,13 @@
 
 sample_size = 3125  # This number is chosen based on the values below.
 
-sample_space = gym.spaces.Box(
+state_sample_space = gym.spaces.Box(
     low=-1.1, high=1.1, shape=env.state_space.shape, dtype=env.state_space.dtype
 )
 
-state_sampler = repeat(
-    grid_sampler(make_grid_from_space(sample_space=sample_space, resolution=25)), num=5
-)
-
-action_sampler = grid_sampler(make_grid_from_ranges([np.linspace(-1, 1, 5)]))
-
-S = sample(
-    sampler=default_sampler(
-        state_sampler=state_sampler, action_sampler=action_sampler, env=env
-    ),
-    sample_size=sample_size,
-)
+state_sampler = grid_sampler(boxgrid(space=state_sample_space, resolution=25)).repeat(5)
+action_sampler = grid_sampler(cartesian(np.linspace(-1, 1, 5)))
+S = transition_sampler(env, state_sampler, action_sampler).sample(size=sample_size)
 
 # %% [markdown]
 #
@@ -91,7 +80,7 @@
 # Generate test points.
 x1 = np.linspace(-1, 1, 50)
 x2 = np.linspace(-1, 1, 50)
-T = make_grid_from_ranges([x1, x2])
+T = cartesian(x1, x2)
 
 # %% [markdown]
 # ## Algorithm
@@ -109,7 +98,7 @@
     target_tube=target_tube,
     problem="FHT",
     regularization_param=regularization_param,
-    kernel_fn=partial(rbf_kernel, gamma=1 / (2 * (sigma ** 2))),
+    kernel_fn=partial(rbf_kernel, sigma=sigma),
     verbose=False,
 )
 
diff --git a/examples/reach/stoch_reach_maximal.py b/examples/reach/stoch_reach_maximal.py
index 14ebe81..5407c18 100644
--- a/examples/reach/stoch_reach_maximal.py
+++ b/examples/reach/stoch_reach_maximal.py
@@ -23,17 +23,15 @@
 from functools import partial
 
 from gym.envs.registration import make
-from sklearn.metrics.pairwise import rbf_kernel
+from gym_socks.kernel.metrics import rbf_kernel
 
 from gym_socks.algorithms.reach.kernel_sr_max import kernel_sr_max
 
-from gym_socks.sampling import sample
-from gym_socks.sampling import default_sampler
+from gym_socks.sampling import transition_sampler
 from gym_socks.sampling import grid_sampler
-from gym_socks.sampling import repeat
 
-from gym_socks.utils.grid import make_grid_from_space
-from gym_socks.utils.grid import make_grid_from_ranges
+from gym_socks.utils.grid import boxgrid
+from gym_socks.utils.grid import cartesian
 
 # %% [markdown]
 # Configuration variables.
@@ -57,25 +55,16 @@
 
 sample_size = 3125  # This number is chosen based on the values below.
 
-sample_space = gym.spaces.Box(
+state_sample_space = gym.spaces.Box(
     low=-1.1, high=1.1, shape=env.state_space.shape, dtype=env.state_space.dtype
 )
 
-state_sampler = repeat(
-    grid_sampler(make_grid_from_space(sample_space=sample_space, resolution=25)), num=5
-)
-
-action_sampler = grid_sampler(make_grid_from_ranges([np.linspace(-1, 1, 5)]))
-
-S = sample(
-    sampler=default_sampler(
-        state_sampler=state_sampler, action_sampler=action_sampler, env=env
-    ),
-    sample_size=sample_size,
-)
+state_sampler = grid_sampler(boxgrid(space=state_sample_space, resolution=25)).repeat(5)
+action_sampler = grid_sampler(cartesian(np.linspace(-1, 1, 5)))
+S = transition_sampler(env, state_sampler, action_sampler).sample(size=sample_size)
 
 # Generate a set of admissible control actions.
-A = make_grid_from_ranges([np.linspace(-1, 1, 5)])
+A = cartesian(np.linspace(-1, 1, 5))
 
 # %% [markdown]
 #
@@ -95,7 +84,7 @@
 # Generate test points.
 x1 = np.linspace(-1, 1, 50)
 x2 = np.linspace(-1, 1, 50)
-T = make_grid_from_ranges([x1, x2])
+T = cartesian(x1, x2)
 
 # %% [markdown]
 # ## Algorithm
@@ -114,7 +103,7 @@
     target_tube=target_tube,
     problem="FHT",
     regularization_param=regularization_param,
-    kernel_fn=partial(rbf_kernel, gamma=1 / (2 * (sigma ** 2))),
+    kernel_fn=partial(rbf_kernel, sigma=sigma),
     verbose=False,
 )
 
diff --git a/gym_socks/algorithms/control/kernel_control_bwd.py b/gym_socks/algorithms/control/kernel_control_bwd.py
index 8b84232..3183873 100644
--- a/gym_socks/algorithms/control/kernel_control_bwd.py
+++ b/gym_socks/algorithms/control/kernel_control_bwd.py
@@ -47,7 +47,7 @@
 
 from gym_socks.sampling.transform import transpose_sample
 
-from gym_socks.utils.batch import generate_batches
+from gym_socks.utils.batch import batch_generator
 
 import logging
 from gym_socks.utils.logging import ms_tqdm, _progress_fmt
@@ -160,7 +160,7 @@ def _compute_backward_recursion_batch(
         # D = np.empty((1, len(Y)))
         D = np.empty_like(CUA)
 
-        for batch in generate_batches(num_elements=len(Y), batch_size=batch_size):
+        for batch in batch_generator(Y, size=batch_size):
 
             beta = np.einsum("ij,ik->ijk", CXY[:, batch], CUA)
             C[batch, :] = np.einsum("i,ijk->jk", Z, beta)
@@ -192,7 +192,7 @@ def _compute_backward_recursion_batch(
 
         #     #     C = np.zeros_like(CUA)
 
-        #     #     for batch in generate_batches(num_elements=len(Y), batch_size=batch_size):
+        #     #     for batch in batch_generator(Y, size=batch_size):
 
         #     #         beta = np.einsum("ij,ik->ijk", CXY[:, batch], CUA)
         #     #         C[batch, :] = np.einsum("i,ijk->jk", Z, beta)
@@ -377,10 +377,9 @@ def train(self, S: np.ndarray, A: np.ndarray):
 
         logger.debug("Computing matrix inverse.")
         self.W = regularized_inverse(
-            X,
-            U=U,
-            kernel_fn=self.kernel_fn,
-            regularization_param=self.regularization_param,
+            self.kernel_fn(X) * self.kernel_fn(U),
+            self.regularization_param,
+            copy=False,
         )
 
         logger.debug("Computing covariance matrices.")
@@ -415,6 +414,8 @@ def __call__(self, time=0, state=None, *args, **kwargs):
             print("Must supply a state to the policy.")
             return None
 
+        state = np.atleast_2d(np.asarray(state, dtype=np.float32))
+
         T = np.array(state)
 
         # Compute covariance matrix.
diff --git a/gym_socks/algorithms/control/kernel_control_fwd.py b/gym_socks/algorithms/control/kernel_control_fwd.py
index e042f3a..16636f0 100644
--- a/gym_socks/algorithms/control/kernel_control_fwd.py
+++ b/gym_socks/algorithms/control/kernel_control_fwd.py
@@ -78,6 +78,7 @@ def kernel_control_fwd(
 
     Returns:
         The policy.
+
     """
 
     alg = KernelControlFwd(
@@ -190,10 +191,9 @@ def train(self, S: np.ndarray, A: np.ndarray):
 
         logger.debug("Computing matrix inverse.")
         self.W = regularized_inverse(
-            X,
-            U=U,
-            kernel_fn=self.kernel_fn,
-            regularization_param=self.regularization_param,
+            self.kernel_fn(X) * self.kernel_fn(U),
+            self.regularization_param,
+            copy=False,
         )
 
         logger.debug("Computing covariance matrix.")
@@ -214,6 +214,8 @@ def __call__(self, time=0, state=None, *args, **kwargs):
             print("Must supply a state to the policy.")
             return None
 
+        state = np.atleast_2d(np.asarray(state, dtype=np.float32))
+
         T = np.array(state)
 
         # Compute covariance matrix.
diff --git a/gym_socks/algorithms/kernel.py b/gym_socks/algorithms/kernel.py
index b815721..4300f8d 100644
--- a/gym_socks/algorithms/kernel.py
+++ b/gym_socks/algorithms/kernel.py
@@ -8,8 +8,6 @@
 
 from gym_socks.algorithms.base import RegressorMixin
 
-from sklearn.preprocessing import normalize
-
 
 def _regression_score(y_true, y_pred):
     return np.abs(np.asarray(y_true) - np.asarray(y_pred)).mean()
@@ -87,4 +85,3 @@ def predict(self, y: np.ndarray, K: np.ndarray):
 
     def score(self, y_test: np.ndarray, K: np.ndarray, y_train: np.ndarray):
         return _regression_score(y_test, self.predict(y_train, K))
-
diff --git a/gym_socks/algorithms/reach/kernel_sr.py b/gym_socks/algorithms/reach/kernel_sr.py
index a9ff19a..c1a6a75 100644
--- a/gym_socks/algorithms/reach/kernel_sr.py
+++ b/gym_socks/algorithms/reach/kernel_sr.py
@@ -20,7 +20,7 @@
 
 from gym_socks.utils import indicator_fn
 from gym_socks.utils import normalize
-from gym_socks.utils.batch import generate_batches
+from gym_socks.utils.batch import batch_generator
 
 import logging
 from gym_socks.utils.logging import ms_tqdm, _progress_fmt
@@ -98,7 +98,7 @@ def _compute_backward_recursion_batch(
     # Backward recursion.
     for t in range(time_horizon - 2, -1, -1):
 
-        for batch in generate_batches(len(Y), batch_size=batch_size):
+        for batch in batch_generator(Y, size=batch_size):
 
             # Compute beta.
             betaXY = normalize(np.einsum("ii,ij->ij", W, CXY[:, batch]))
@@ -287,7 +287,9 @@ def fit(self, S: np.ndarray):
 
         logger.debug("Computing matrix inverse.")
         self.W = regularized_inverse(
-            X, kernel_fn=self.kernel_fn, regularization_param=self.regularization_param
+            self.kernel_fn(X),
+            self.regularization_param,
+            copy=False,
         )
 
         logger.debug("Computing covariance matrix.")
diff --git a/gym_socks/algorithms/reach/kernel_sr_max.py b/gym_socks/algorithms/reach/kernel_sr_max.py
index ed8e57b..0336083 100644
--- a/gym_socks/algorithms/reach/kernel_sr_max.py
+++ b/gym_socks/algorithms/reach/kernel_sr_max.py
@@ -19,7 +19,7 @@
 
 from gym_socks.utils import indicator_fn
 from gym_socks.utils import normalize
-from gym_socks.utils.batch import generate_batches
+from gym_socks.utils.batch import batch_generator
 
 import logging
 from gym_socks.utils.logging import ms_tqdm, _progress_fmt
@@ -109,7 +109,7 @@ def _compute_backward_recursion_batch(
     # Backward recursion.
     for t in range(time_horizon - 2, -1, -1):
 
-        for batch in generate_batches(len(Y), batch_size=batch_size):
+        for batch in batch_generator(Y, size=batch_size):
 
             # Compute beta.
             betaXY = normalize(
@@ -321,10 +321,9 @@ def fit(self, S: np.ndarray, A: np.ndarray):
 
         logger.debug("Computing matrix inverse.")
         self.W = regularized_inverse(
-            X,
-            U=U,
-            kernel_fn=self.kernel_fn,
-            regularization_param=self.regularization_param,
+            self.kernel_fn(X) * self.kernel_fn(U),
+            self.regularization_param,
+            copy=False,
         )
 
         logger.debug("Computing covariance matrices.")
diff --git a/gym_socks/algorithms/reach/monte_carlo.py b/gym_socks/algorithms/reach/monte_carlo.py
index fba14c3..de622a2 100644
--- a/gym_socks/algorithms/reach/monte_carlo.py
+++ b/gym_socks/algorithms/reach/monte_carlo.py
@@ -1,3 +1,5 @@
+from gym_socks.sampling import sample_fn
+
 """Stochastic reachability using Monte-Carlo."""
 
 import numpy as np
@@ -8,8 +10,7 @@
 from gym_socks.envs.dynamical_system import DynamicalSystem
 from gym_socks.policies import BasePolicy
 
-from gym_socks.sampling import sample
-from gym_socks.sampling import sample_generator
+from gym_socks.sampling import sample_fn
 
 from gym_socks.utils import indicator_fn
 
@@ -36,6 +37,7 @@ def _trajectory_indicator(
     return result
 
 
+@sample_fn
 def _monte_carlo_trajectory_sampler(
     time_horizon: int = None,
     env: DynamicalSystem = None,
@@ -54,26 +56,21 @@ def _monte_carlo_trajectory_sampler(
 
     """
 
-    @sample_generator
-    def _sample_generator():
-
-        state_sequence = []
-        state_sequence.append(state)
-
-        env.state = state
+    state_sequence = []
+    state_sequence.append(state)
 
-        time = 0
-        for t in range(time_horizon):
-            action = policy(time=time, state=env.state)
-            next_state, cost, done, _ = env.step(time=t, action=action)
+    env.reset(state)
 
-            state_sequence.append(next_state)
+    time = 0
+    for t in range(time_horizon):
+        action = policy(time=time, state=env.state)
+        next_state, *_ = env.step(time=t, action=action)
 
-            time += 1
+        state_sequence.append(next_state)
 
-        yield state_sequence
+        time += 1
 
-    return _sample_generator
+    yield state_sequence
 
 
 def monte_carlo_sr(
@@ -236,12 +233,9 @@ def fit_predict(self, env: DynamicalSystem, policy: BasePolicy, T: np.ndarray):
         for i, state in enumerate(T):
 
             # For each point, generate a collection of trajectories.
-            S = sample(
-                sampler=_monte_carlo_trajectory_sampler(
-                    time_horizon=self.time_horizon, env=env, state=state, policy=policy
-                ),
-                sample_size=self.num_iterations,
-            )
+            S = _monte_carlo_trajectory_sampler(
+                time_horizon=self.time_horizon, env=env, state=state, policy=policy
+            ).sample(size=self.num_iterations)
 
             # Working backwards in time, compute the "likelihood" that the trajectories
             # will satisfy the constraints set up by the constraint tube and target
diff --git a/gym_socks/algorithms/reach/separating_kernel.py b/gym_socks/algorithms/reach/separating_kernel.py
index ea08b9d..f6db607 100644
--- a/gym_socks/algorithms/reach/separating_kernel.py
+++ b/gym_socks/algorithms/reach/separating_kernel.py
@@ -6,27 +6,27 @@
 
 from functools import partial
 
+import numpy as np
+
 from gym_socks.algorithms.base import ClassifierMixin
 from gym_socks.kernel.metrics import abel_kernel
 from gym_socks.kernel.metrics import regularized_inverse
 
-import numpy as np
+from gym_socks.utils.validation import check_array
 
 
 class SeparatingKernelClassifier(ClassifierMixin):
-    """Separating kernel classifier.
+    r"""Separating kernel classifier.
 
-    A kernel-based support classifier for unknown distributions. Given a set of data taken iid from the distribution, the `SeparatingKernelClassifier` constructs a kernel-based classifier of the support of the distribution based on the theory of separating kernels.
-
-    Note:
-        The sample used by the classifier is from the marginal distribution, not the
-        joint or conditional. Thus, the data should be an array of points organized such
-        that each point occupies a single row in a 2D-array.
+    A kernel-based support classifier for unknown distributions. Given a set of data
+    taken i.i.d. from the distribution, the `SeparatingKernelClassifier` constructs a
+    kernel-based classifier of the support of the distribution based on the theory of
+    separating kernels.
 
     Args:
-        kernel_fn: The kernel function used by the classifier.
-        regularization_param: The regularization parameter used in the regularized
-            least-squares problem. Determines the smoothness of the solution.
+        kernel_fn: The kernel function used by the algorithm.
+        regularization_param: The regularization parameter :math:`\lambda > 0`.
+            Determines the smoothness of the solution.
 
     Example:
         >>> from gym_socks.algorithms.reach import SeparatingKernelClassifier
@@ -42,7 +42,7 @@ class SeparatingKernelClassifier(ClassifierMixin):
     def __init__(
         self,
         kernel_fn=None,
-        regularization_param=None,
+        regularization_param: float = None,
         *args,
         **kwargs,
     ):
@@ -51,22 +51,15 @@ def __init__(
         self.kernel_fn = kernel_fn
         self.regularization_param = regularization_param
 
-    def _validate_params(self, S):
-
-        if self.kernel_fn is None:
-            self.kernel_fn = partial(abel_kernel, sigma=0.1)
-
-        if self.regularization_param is None:
-            self.regularization_param = 1
-
-    def _validate_data(self, X):
-
-        if X is None:
-            raise ValueError("Must supply a sample.")
-
     def fit(self, X: np.ndarray):
         """Fit separating kernel classifier.
 
+        Note:
+
+            The sample used by the classifier is from the marginal distribution, not the
+            joint or conditional. Thus, the data should be an array of points organized
+            such that each point occupies a single row in a 2D-array.
+
         Args:
             X: Data drawn from distribution.
 
@@ -75,18 +68,19 @@ def fit(self, X: np.ndarray):
 
         """
 
-        self._validate_data(X)
-        X = np.array(X)
+        if self.kernel_fn is None:
+            self.kernel_fn = partial(abel_kernel, sigma=0.1)
+
+        self._X = check_array(X)
 
-        self.X = X
+        if self.regularization_param is None:
+            self.regularization_param = 1 / len(self._X)
 
+        # Precompute matrices which are used in prediction.
         K = self.kernel_fn(X)
+        self._W = regularized_inverse(K, self.regularization_param, copy=True)
 
-        self.W = regularized_inverse(
-            X, kernel_fn=self.kernel_fn, regularization_param=self.regularization_param
-        )
-
-        self.tau = 1 - np.min(np.diagonal((1 / len(X)) * K.T @ self.W @ K))
+        self._tau = 1 - np.min(np.diagonal(K.T @ self._W @ K))
 
         return self
 
@@ -101,14 +95,12 @@ def predict(self, T: np.ndarray) -> list:
 
         """
 
-        self._validate_data(T)
-        T = np.array(T)
-
-        K = self.kernel_fn(self.X, T)
+        T = check_array(T)
 
-        C = np.diagonal((1 / len(self.X)) * K.T @ self.W @ K)
+        K = self.kernel_fn(self._X, T)
+        C = np.diagonal(K.T @ self._W @ K)
 
-        return C >= 1 - self.tau
+        return np.array(C >= 1 - self._tau, dtype=bool)
 
     def score(self):
         raise NotImplementedError
diff --git a/gym_socks/envs/__init__.py b/gym_socks/envs/__init__.py
index fbbf109..988cdb6 100644
--- a/gym_socks/envs/__init__.py
+++ b/gym_socks/envs/__init__.py
@@ -1,13 +1,11 @@
 __all__ = [
-    # Systems
     "cwh",
     "integrator",
     "nonholonomic",
+    "planar_quad",
     "point_mass",
     "QUAD20",
     "tora",
-    # Classes
-    "policy",
 ]
 
 from gym_socks.envs.cwh import CWH4DEnv
diff --git a/gym_socks/envs/core.py b/gym_socks/envs/core.py
index 4bd1066..145adca 100644
--- a/gym_socks/envs/core.py
+++ b/gym_socks/envs/core.py
@@ -53,7 +53,7 @@ def step(self):
         raise NotImplementedError
 
     @abstractmethod
-    def reset(self):
+    def reset(self, state=None):
         """Reset the object to an initial state.
 
         Resets the object to an initial state (which may be random).
@@ -137,42 +137,3 @@ def __init__(self, env):
         self.__dict__ = env.__dict__
 
         self.env = env
-
-
-# class HookWrapper(BaseWrapper):
-#     def __getattr__(self, attr):
-#         """Wrapped __getattr__."""
-#         _attr = self.env.__getattribute__(attr)
-
-#         if callable(_attr):
-
-#             @wraps(callable)
-#             def _wrapper(*args, **kwargs):
-#                 self.pre_hook()
-#                 result = _attr(*args, **kwargs)
-#                 self.post_hook()
-
-#                 # prevent unwrapping
-#                 if result is self.env:
-#                     return self
-
-#                 return result
-
-#             return _wrapper
-
-#         else:
-#             return _attr
-
-#     # @abstractmethod
-#     def pre_hook(self):
-#         raise NotImplementedError
-
-#     # @abstractmethod
-#     def post_hook(self):
-#         raise NotImplementedError
-
-#     # def reset(self):
-#     #     self.pre_hook()
-#     #     result = self.env.reset()
-#     #     self.post_hook()
-#     #     return result
diff --git a/gym_socks/envs/dynamical_system.py b/gym_socks/envs/dynamical_system.py
index d8b0784..b92e289 100644
--- a/gym_socks/envs/dynamical_system.py
+++ b/gym_socks/envs/dynamical_system.py
@@ -222,7 +222,7 @@ def generate_observation(self, time, state, action):
 
         """
 
-        return np.array(state, dtype=np.float32)
+        return np.asarray(state, dtype=np.float32)
 
     def cost(self, time, state, action):
         """Cost function for the system.
@@ -257,7 +257,7 @@ def step(self, time=0, action=None) -> tuple:
 
         """
 
-        action = np.asarray(action, dtype=np.float32)
+        action = np.atleast_1d(np.asarray(action, dtype=np.float32))
 
         err_msg = "%r (%s) invalid" % (action, type(action))
         assert self.action_space.contains(action), err_msg
@@ -293,10 +293,14 @@ def step(self, time=0, action=None) -> tuple:
 
         return observation, cost, done, info
 
-    def reset(self):
+    def reset(self, state=None):
         """Reset the system to a random initial condition."""
-        self.state = self.state_space.sample()
-        return np.array(self.state, dtype=np.float32)
+        if state is None:
+            self.state = self.state_space.sample()
+        else:
+            self.state = np.atleast_1d(np.asarray(state, dtype=np.float32))
+
+        return self.state
 
     def render(self, mode="human"):
         raise NotImplementedError
@@ -315,3 +319,42 @@ def np_random(self):
     def seed(self, seed=None):
         self._np_random, seed = seeding.np_random(seed)
         return [seed]
+
+
+def simulate(env, N, x0=None, policy=None, full_sequence=False):
+    """Helper function to simulate a dynamical system.
+
+    Args:
+        env: Environment to simulate.
+        N: The time horizon for the simulation.
+        x0: The initial condition of the system.
+        policy: The policy.
+
+    Returns:
+        A tuple with 2 elements containing the state sequence and the action sequence.
+
+    """
+
+    state_sequence = []
+    action_sequence = []
+
+    env.reset(x0)
+
+    if full_sequence is True:
+        state_sequence.append(env.state)
+
+    for t in range(N):
+        if policy is not None:
+            action = policy(time=t, state=env.state)
+        else:
+            action = env.action_space.sample()
+
+        action_sequence.append(action)
+
+        env.step(time=t, action=action)
+        state_sequence.append(env.state)
+
+    return (
+        np.array(state_sequence, dtype=np.float32),
+        np.array(action_sequence, dtype=np.float32),
+    )
diff --git a/gym_socks/envs/integrator.py b/gym_socks/envs/integrator.py
index 02f6c7f..78fedb1 100644
--- a/gym_socks/envs/integrator.py
+++ b/gym_socks/envs/integrator.py
@@ -118,8 +118,6 @@ class NDIntegratorEnv(DynamicalSystem):
         >>> for i in range(10):
         ...     action = env.action_space.sample()
         ...     obs, reward, done, _ = env.step(action)
-        ...
-        >>> env.close()
 
     """
 
diff --git a/gym_socks/envs/planar_quad.py b/gym_socks/envs/planar_quad.py
index eed0496..9749ad9 100644
--- a/gym_socks/envs/planar_quad.py
+++ b/gym_socks/envs/planar_quad.py
@@ -14,7 +14,10 @@ class PlanarQuadrotorEnv(DynamicalSystem):
 
     Bases: :py:class:`gym_socks.envs.dynamical_system.DynamicalSystem`
 
-    A planar quadrotor is quadrotor restricted to two dimensions. Similar to the OpenAI gym lunar lander benchmark, the planar quadrotor is a bar with two independent rotors at either end. Inputs are the trust of the rotors, and apply a torque to the bar. The system is also subject to gravitational forces.
+    A planar quadrotor is quadrotor restricted to two dimensions. Similar to the OpenAI
+    gym lunar lander benchmark, the planar quadrotor is a bar with two independent
+    rotors at either end. Inputs are the trust of the rotors, and apply a torque to the
+    bar. The system is also subject to gravitational forces.
 
     """
 
diff --git a/gym_socks/kernel/metrics.py b/gym_socks/kernel/metrics.py
index 374be7e..9b08f08 100644
--- a/gym_socks/kernel/metrics.py
+++ b/gym_socks/kernel/metrics.py
@@ -2,164 +2,228 @@
 Kernel functions and helper utilities for kernel-based calculations.
 
 Most of the commonly-used kernel functions are already implemented in
-sklearn.metrics.pairwise. The RBF kernel and pairwise Euclidean distance function is
-re-implemented here as an alternative, in case sklearn is unavailable. Most, if not all
-of the kernel functions defined in sklearn.metrics.pairwise should be compatible with
-the functions defined here.
+:py:mod:`sklearn.metrics.pairwise`. The RBF kernel and pairwise Euclidean distance
+function is re-implemented here as an alternative, in case :py:mod:`sklearn` is
+unavailable. Most, if not all of the kernel functions defined in
+:py:mod:`sklearn.metrics.pairwise` should be compatible with the functions defined here.
 
-To use the sklearn functions, use them like so:
+Attention:
 
-    >>> from sklearn.metrics.pairwise import rbf_kernel
-    >>> from sklearn.metrics.pairwise import euclidean_distances
-    >>> X = np.arange(4).reshape((2, 2))
-    >>> Y = np.arange(6).reshape((3, 2))
-    >>> K = gym_socks.kernel.metrics.rbf_kernel(X, Y, distance_fn=euclidean_distances)
+    In Matlab, data is typically ordered differently than in Python. In Matlab, data is
+    ordered in columns, whereas in Python, data is ordered in **rows**. If you import
+    data from a Matlab file, be sure to transpose it if needed to follow Python
+    formatting, i.e. ``X = X.T``.
+
+    For example, the data in ``X`` and ``Y`` should be organized as::
+
+        X = [[--- x1 ---],
+             [--- x2 ---],
+             ...
+             [--- xn ---]]
+
+        Y = [[--- y1 ---],
+             [--- y2 ---],
+             ...
+             [--- ym ---]]
 
 """
 
 from functools import partial, reduce
 
 import numpy as np
-from numpy.linalg import inv
 
+from gym_socks.utils.validation import check_array
+from gym_socks.utils.validation import check_matrix
 
-def euclidean_distance(
-    X: np.ndarray,
-    Y: np.ndarray = None,
-    squared: bool = False,
-) -> np.ndarray:
-    """Euclidean distance function.
 
-    The main difference between the way this function calculates euclidean
-    distance over other implementations such as in sklearn.metrics.pairwise is
-    that this implementation is largely agnostic to the dimensionality of the
-    input data, and generally works well for high-dimensional vectors and dense
-    data sets, such as state trajectories over a long time horizon.
+def check_pairwise_arrays(
+    X,
+    Y=None,
+    dtype: np.dtype = np.float64,
+    ensure_finite: bool = True,
+    copy: bool = False,
+):
+    """Check pairwise arrays.
+
+    Note:
+
+        This function is intended as a simple replacement for the
+        :py:func:`sklearn.metrics.pairwise.check_pairwise_arrays` function. Unlike
+        :py:mod:`sklearn`, this function does not check sparse input data, and does not
+        do sophisticated type checking or upcasting.
 
     Args:
-        X: The observations oganized in ROWS.
-        Y: The observations oganized in ROWS.
-        squared: Whether or not the result is the squared Euclidean distance.
+        X: A 2D array with observations oganized in ROWS.
+        Y: A 2D array with observations oganized in ROWS.
+        dtype: The data type of the resulting array.
+        copy: Whether to create a forced copy of ``array``.
+        ensure_finite: Whether to raise an error if the array is not finite.
 
     Returns:
-        np.ndarray: Matrix of pairwise distances between points.
+        The validated arrays ``X`` and ``Y``.
 
     """
 
     if Y is None:
+        X = check_array(X, dtype=dtype, ensure_finite=ensure_finite, copy=copy)
         Y = X
 
-    num_rows_X, num_cols_X = X.shape
-    num_rows_Y, num_cols_Y = Y.shape
+    else:
+        X = check_array(X, dtype=dtype, ensure_finite=ensure_finite, copy=copy)
+        Y = check_array(Y, dtype=dtype, ensure_finite=ensure_finite, copy=copy)
+
+    return X, Y
 
-    def calc_dim(prev, next):
-        x, y = next
-        return prev + np.power(
-            np.tile(y, (num_rows_X, 1)) - np.tile(x, (num_rows_Y, 1)).T,
-            2,
-        )
 
-    distances = reduce(
-        calc_dim, zip(X.T, Y.T), np.zeros(shape=(num_rows_X, num_rows_Y))
-    )
+def euclidean_distances(X, Y=None, squared: bool = False):
+    """Compute the pairwise Euclidean distance matrix between points.
 
-    return distances if squared else np.sqrt(distances)
+    Note:
+
+        This function is intended as a simple replacement for the
+        :py:func:`sklearn.metrics.pairwise.euclidean_distances` function. Unlike
+        :py:mod:`sklearn`, this function does not check sparse input data, and does not
+        do sophisticated type checking or upcasting.
+
+    Args:
+        X: A 2D array with observations oganized in ROWS.
+        Y: A 2D array with observations oganized in ROWS.
+        squared: Whether the result is squared before returning.
+
+    Returns:
+        The matrix of pairwise Euclidean distances between points.
+
+    """
+
+    X, Y = check_pairwise_arrays(X, Y)
+
+    XX = np.einsum("ij,ij->i", X, X)[:, np.newaxis]
+    YY = np.einsum("ij,ij->i", Y, Y)[np.newaxis, :]
+
+    distances = -2 * (X @ Y.T)
+    distances += XX
+    distances += YY
+
+    np.maximum(distances, 0, out=distances)
+    if X is Y:
+        np.fill_diagonal(distances, 0)
+
+    return distances if squared else np.sqrt(distances, out=distances)
+
+
+def check_pairwise_distances(D: np.ndarray, shape: tuple, copy: bool = True):
+    """Validate the pairwise distance matrix.
+
+    Performs checks to ensure the pairwise distance matrix is valid.
+
+    Args:
+        D: Pairwise distance matrix.
+        shape: The desired shape of the array.
+        copy: Whether to create a forced copy of ``D``.
+
+    Returns:
+        The validated matrix.
+
+    """
+
+    D = check_array(D, copy=copy)
+
+    if D.shape != shape:
+        raise ValueError(
+            f"Expected pairwise distances to be of shape {shape}, "
+            f"but got {D.shape} instead."
+        )
+
+    return D
 
 
 def rbf_kernel(
     X: np.ndarray,
     Y: np.ndarray = None,
     sigma: float = None,
-    distance_fn=None,
+    D: np.ndarray = None,
 ) -> np.ndarray:
-    """RBF kernel function.
+    r"""RBF kernel function.
 
-    Computes the pairwise evaluation of the RBF kernel on each vector in X and
-    Y. For instance, if X has n vecotrs, and Y has m vectors, then the result is
-    an n x m matrix K where K_ij = k(xi, yj).
-    ::
+    Computes the pairwise evaluation of the RBF kernel on each vector in ``X`` and
+    ``Y``. For example, if ``X`` has :math:`m` vecotrs, and ``Y`` has :math:`n` vectors,
+    then the result is an :math:`m \times n` matrix :math:`K` where :math:`K_{ij} =
+    k(x_i, y_j)`.
 
-        X = [[--- x1 ---],
-             [--- x2 ---],
-             ...
-             [--- xn ---]]
+    .. math::
 
-        Y = [[--- y1 ---],
-             [--- y2 ---],
-             ...
-             [--- ym ---]]
+        K =
+        \begin{bmatrix}
+            k(x_1,y_1) & \cdots & k(x_1,y_n) \\
+            \vdots & \ddots & \vdots \\
+            k(x_m,y_1) & \cdots & k(x_m,y_n)
+        \end{bmatrix}
 
-    The result is a matrix,
-    ::
+    Note:
 
-        K = [[k(x1,y1), ..., k(x1,ym)],
-             ...
-             [k(xn,y1), ..., k(xn,ym)]]
+        This function is intended as a simple replacement for the
+        :py:func:`sklearn.metrics.pairwise.rbf_kernel` function. Unlike
+        :py:mod:`sklearn`, this function does not check sparse input data, and does not
+        do sophisticated type checking or upcasting.
 
-    The main difference between this implementation and the rbf_kernel in
-    sklearn.metrics.pairwise is that this function optionally allows you to
-    specify a different distance metric in the event the data is non-Euclidean.
+        The main difference between this implementation and the
+        :py:func:`sklearn.metrics.pairwise.rbf_kernel` in
+        :py:mod:`sklearn.metrics.pairwise` is that this function optionally allows you
+        to specify a different distance metric in the event the data is non-Euclidean.
 
-    Args:
-        X: The observations oganized in ROWS.
-        Y: The observations oganized in ROWS.
-        sigma: Strictly positive real-valued kernel parameter.
-        distance_fn: Distance function to use in the kernel evaluation.
+    Attention:
 
-    Returns:
-        np.ndarray: Gram matrix of pairwise evaluations of the kernel function.
+        If you are familiar with :py:func:`sklearn.metrics.pairwise.rbf_kernel`, note
+        that the parameter :py:attr:`sigma` is *not* the same as the ``gamma`` parameter
+        used by :py:func:`sklearn.metrics.pairwise.rbf_kernel`. However, they are
+        related:
 
-    """
+        .. math::
 
-    # if distance_fn is None:
-    #     distance_fn = euclidean_distance
+            \gamma = \frac{1}{2 \sigma^{2}}
 
-    # K = distance_fn(X, Y, squared=True)
+    Args:
+        X: A 2D array with observations oganized in ROWS.
+        Y: A 2D array with observations oganized in ROWS.
+        sigma: Strictly positive real-valued kernel parameter.
+        D: Pairwise distance matrix.
 
-    if Y is None:
-        Y = X
+    Returns:
+        Gram matrix of pairwise evaluations of the kernel function.
 
-    N, M = np.shape(X)
-    T = len(Y)
+    """
 
-    K = np.zeros((N, T))
+    X, Y = check_pairwise_arrays(X, Y)
 
-    for i in range(M):
-        K += np.power(np.tile(Y[:, i], (N, 1)) - np.tile(X[:, i], (T, 1)).T, 2)
+    if D is None:
+        D = euclidean_distances(X, Y, squared=True)
+    else:
+        D = check_pairwise_distances(D, shape=(X.shape[0], Y.shape[0]))
 
     if sigma is None:
-        sigma = np.median(K)
+        sigma = np.median(D)
     else:
         assert sigma > 0, "sigma must be a strictly positive real value."
 
-    K /= -2 * (sigma ** 2)
-    np.exp(K, K)
+    D /= -2 * (sigma ** 2)
+    np.exp(D, D)
 
-    return K
+    return D
 
 
 def rbf_kernel_derivative(
     X: np.ndarray,
     Y: np.ndarray = None,
     sigma: float = None,
-    distance_fn=None,
+    D: np.ndarray = None,
 ) -> np.ndarray:
 
-    # if distance_fn is None:
-    #     distance_fn = euclidean_distance
-
-    # D = distance_fn(X, Y, squared=False)
-
-    if Y is None:
-        Y = X
-
-    N, M = np.shape(X)
-    T = len(Y)
-
-    D = np.zeros((N, T))
+    X, Y = check_pairwise_arrays(X, Y)
 
-    for i in range(M):
-        D += np.tile(Y[:, i], (N, 1)) - np.tile(X[:, i], (T, 1)).T
+    if D is None:
+        D = euclidean_distances(X, Y, squared=False)
+    else:
+        D = check_pairwise_distances(D, shape=(X.shape[0], Y.shape[0]))
 
     if sigma is None:
         sigma = np.median(D)
@@ -175,93 +239,142 @@ def abel_kernel(
     X: np.ndarray,
     Y: np.ndarray = None,
     sigma: float = None,
-    distance_fn=None,
+    D: np.ndarray = None,
 ) -> np.ndarray:
     """Abel kernel function.
 
     Args:
-        X: The observations oganized in ROWS.
-        Y: The observations oganized in ROWS.
+        X: A 2D array with observations oganized in ROWS.
+        Y: A 2D array with observations oganized in ROWS.
         sigma: Strictly positive real-valued kernel parameter.
-        distance_fn: Distance function to use in the kernel evaluation.
+        D: Pairwise distance matrix.
 
     Returns:
-        np.ndarray: Gram matrix of pairwise evaluations of the kernel function.
+        Gram matrix of pairwise evaluations of the kernel function.
 
     """
 
-    if distance_fn is None:
-        distance_fn = euclidean_distance
+    X, Y = check_pairwise_arrays(X, Y)
 
-    K = distance_fn(X, Y, squared=False)
+    if D is None:
+        D = euclidean_distances(X, Y, squared=False)
+    else:
+        D = check_pairwise_distances(D, shape=(X.shape[0], Y.shape[0]))
 
     if sigma is None:
-        sigma = np.median(K)
+        sigma = np.median(D)
     else:
         assert sigma > 0, "sigma must be a strictly positive real value."
 
-    K /= -sigma
-    np.exp(K, K)
+    D /= -sigma
+    np.exp(D, D)
 
-    return K
+    return D
 
 
 def delta_kernel(X: np.ndarray, Y: np.ndarray = None):
-    """Delta (discrete) kernel function.
+    r"""Delta (discrete) kernel function.
 
     The delta kernel is defined as :math:`k(x_{i}, x_{j}) = \delta_{ij}`, meaning the
     kernel returns a 1 if the vectors are the same, and a 0 otherwise. The vectors in
-    `X` and `Y` should have discrete values, meaning each element in the vector should
+    ``X`` and ``Y`` should have discrete values, meaning each element in the vector should
     be a natural number or integer value.
 
     Args:
-        X: A 2D ndarray with the observations oganized in ROWS.
-        Y: A 2D ndarray with the observations oganized in ROWS.
+        X: A 2D array with observations oganized in ROWS.
+        Y: A 2D array with observations oganized in ROWS.
 
     Returns:
-        np.ndarray: Gram matrix of pairwise evaluations of the kernel function.
+        Gram matrix of pairwise evaluations of the kernel function.
 
     """
 
-    if X.ndim == 1:
-        raise ValueError(
-            "Expected 2D array for X, got 1D array instead. \n"
-            "Reshape the data using array.reshape(-1, 1) "
-            "if the data has only a single dimension "
-            "or array.reshape(1, -1) if there is only a single sample."
-        )
+    X, Y = check_pairwise_arrays(X, Y, dtype=int)
 
-    if Y is None:
-        Y = X
+    D = np.all(X[:, :, None] == Y.T, axis=1).astype(int)
+
+    return D
 
-    else:
 
-        if Y.ndim == 1:
-            raise ValueError(
-                "Expected 2D array for Y, got 1D array instead. \n"
-                "Reshape the data using array.reshape(-1, 1) "
-                "if the data has only a single dimension "
-                "or array.reshape(1, -1) if there is only a single sample."
-            )
+def _hybrid_distances(
+    X: np.ndarray, Q: np.ndarray, Y: np.ndarray = None, R: np.ndarray = None
+):
+    r"""Hybrid distance function.
 
-    N, M = np.shape(X)
-    T = len(Y)
+    The hybrid distance function is a distance metric that works for states that are a
+    combination of a continuous state and a discrete mode.
+
+    .. math::
 
-    D = np.ones((N, T), dtype=int)
+        d((x, q), (x', q')) =
+        \begin{cases}
+            \zeta(x - x'), & q = q' \\
+		    1, & q \neq q'
+        \end{cases}
 
-    for i in range(M):
-        D &= np.tile(Y[:, i], (N, 1)) == np.tile(X[:, i], (T, 1)).T
+    .. math::
+
+        \zeta(x) = (2/\pi) \max_{1 \leq i \leq n} \tan^{-1} \vert x_{i} \vert
+
+    """
+
+    X, Y = check_pairwise_arrays(X, Y)
+
+    d = delta_kernel(Q, R)
+
+    diff = np.abs(X[:, :, None] - Y.T)
+    diff = (2 / np.pi) * np.max(np.arctan(diff), axis=1)
+
+    D = np.where(d == 0, 1, diff)
 
     return D
 
 
+def hybrid_kernel(
+    X: np.ndarray, Q: np.ndarray, Y: np.ndarray = None, R: np.ndarray = None
+):
+    r"""Hybrid systems kernel.
+
+    In a hybrid system, we split the sample according to the mode ``Q``. The vectors in
+    ``Q`` and ``R`` should have discrete values, meaning each element in the vector
+    should be a natural number or integer value.
+
+    .. math::
+
+        d((x, q), (x', q')) =
+        \begin{cases}
+            \zeta(x - x'), & q = q' \\
+		    1, & q \neq q'
+        \end{cases}
+
+    .. math::
+
+        \zeta(x) = (2/\pi) \max_{1 \leq i \leq n} \tan^{-1} \vert x_{i} \vert
+
+    .. math::
+
+        k(x, q, x', q') = 1 - d((x, q), (x', q'))
+
+    Args:
+        X: A 2D array with observations oganized in ROWS.
+        Q: A 2D array with observations oganized in ROWS.
+        Y: A 2D array with observations oganized in ROWS.
+        R: A 2D array with observations oganized in ROWS.
+
+    Returns:
+        Gram matrix of pairwise evaluations of the kernel function.
+
+    """
+
+    D = _hybrid_distances(X, Q, Y, R)
+
+    return 1 - np.sqrt(D)
+
+
 def regularized_inverse(
-    X: np.ndarray,
-    Y: np.ndarray = None,
-    U: np.ndarray = None,
-    V: np.ndarray = None,
+    G: np.ndarray,
     regularization_param: float = None,
-    kernel_fn=None,
+    copy: bool = True,
 ) -> np.ndarray:
     r"""Regularized inverse.
 
@@ -269,61 +382,113 @@ def regularized_inverse(
 
     .. math::
 
-        W = (K + \lambda M I)^{-1}, \quad
-        K \in \mathbb{R}^{n \times n}, \quad
-        K_{ij} = k(x_{i}, y_{j})
-
-    Example:
-
-        >>> from gym_socks.kernel.metrics import regularized_inverse
-        >>> X = np.arange(4).reshape((2, 2))
-        >>> Y = np.arange(6).reshape((3, 2))
-        >>> K = regularized_inverse(X, Y)
+        W = (G + \lambda M I)^{-1}, \quad
+        G \in \mathbb{R}^{n \times n}, \quad
+        G_{ij} = k(x_{i}, y_{j})
 
     Args:
-        X: The observations oganized in ROWS.
-        Y: The observations oganized in ROWS.
-        regularization_param: Regularization parameter, which is a strictly positive
-            real value.
-        kernel_fn: The kernel function is a function that returns an ndarray where
-            each element is the pairwise evaluation of a kernel function. See sklearn.
-            metrics.pairwise for more info. The default is the RBF kernel.
+        G: The Gram (kernel) matrix.
+        regularization_param: The regularization parameter :math:`\lambda > 0`.
+        copy: Whether to create a forced copy of ``G``.
 
     Returns:
         Regularized matrix inverse.
 
     """
 
-    if Y is None:
-        Y = X
+    G = check_matrix(G, ensure_square=True, copy=copy)
 
     if regularization_param is None:
-        regularization_param = 1 / (X.shape[0] ** 2)
+        regularization_param = 1 / (len(G) ** 2)
     else:
         assert (
             regularization_param > 0
         ), "regularization_param must be a strictly positive real value."
 
-    if kernel_fn is None:
-        kernel_fn = partial(rbf_kernel, sigma=1)
+    G[np.diag_indices_from(G)] += regularization_param * len(G)
+
+    return np.linalg.inv(G)
+
+
+def woodbury_inverse(
+    A: np.ndarray,
+    U: np.ndarray,
+    C: np.ndarray,
+    V: np.ndarray,
+    precomputed: bool = False,
+) -> np.ndarray:
+    r"""Computes the matrix inverse using the Woodbury matrix identity.
+
+    .. math::
+
+        W = (A + U C V)^{-1} = A^{-1} - A^{-1} U (C^{-1} + V A^{-1} U)^{-1} V A^{-1}
+
+    where :math:`A` is :math:`n \times n`, :math:`C` is :math:`k \times k`, :math:`U` is
+    :math:`n \times k`, and :math:`V` is :math:`k \times n`.
+
+    This function is useful for computing the regularized inverse in a more
+    computationally efficient manner. This happens because the matrices :math:`A` and
+    :math:`C` are typically easy to invert manually, either because the inverses are
+    known a priori, are scalar matrices, or identity, leading to a smaller matrix
+    inversion in the calculations.
+
+    Example:
+
+        >>> import numpy as np
+        >>> from gym_socks.kernel.metrics import regularized_inverse
+        >>> from gym_socks.kernel.metrics import woodbury_inverse
+        >>> X = np.random.randn(100, 2)
+        >>> G = X @ X.T
+        >>> timeit regularized_inverse(G, 1)
+        184 µs ± 7.14 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
+        >>> A = (1 / 100) * np.identity(100)
+        >>> C = np.identity(2)
+        >>> timeit woodbury_inverse(A, X, C, X.T, precomputed=True)
+        78.9 µs ± 1.67 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
+        >>> W1 = regularized_inverse(G, 1)
+        >>> W2 = woodbury_inverse(A, X, C, X.T, precomputed=True)
+        >>> np.allclose(W1, W2)
+        True
+
+    Args:
+        A: A conformable square matrix. Must be nonsingular.
+        U: A conformable matrix.
+        C: A conformable square matrix. Must be nonsingular.
+        V: A conformable matrix.
+        precomputed: Whether ``A`` and ``C`` are the precomputed inverses.
 
-    err_msg = "Parameters %r, %r must have the same shape." % (X, Y)
-    assert X.shape == Y.shape, err_msg
+    Returns:
+        The resulting inverse matrix :math:`W`.
 
-    num_rows, num_cols = X.shape
+    """
 
-    K = kernel_fn(X, Y)
+    A = check_matrix(A, ensure_square=True)
+    C = check_matrix(C, ensure_square=True)
 
-    if U is not None:
-        err_msg = "Parameters %r, %r must have the same sample size." % (X, U)
-        assert X.shape[0] == U.shape[0], err_msg
+    U = check_matrix(U)
+    V = check_matrix(V)
 
-        if V is not None:
-            err_msg = "Parameters %r, %r must have the same shape." % (U, V)
-            assert U.shape == V.shape, err_msg
+    m, n = np.shape(A)
+    k, l = np.shape(C)
 
-        K = np.multiply(kernel_fn(U, V), K)
+    assert U.shape == (n, k), f"Expected U to be {(n, k)}, instead got: {U.shape}"
+    assert V.shape == (k, n), f"Expected V to be {(k, n)}, instead got: {V.shape}"
 
-    I = np.identity(num_rows)
+    if precomputed is False:
+        A_inv = np.linalg.inv(A)
+        C_inv = np.linalg.inv(C)
+    else:
+        A_inv = A
+        C_inv = C
+
+    VA = V @ A_inv
+    D = C_inv + VA @ U
+    if np.isscalar(D):
+        np.reciprocal(D, out=D)
+
+        W = A_inv - D * A_inv @ U @ VA
+
+    else:
+        W = A_inv - A_inv @ U @ np.linalg.solve(D, VA)
 
-    return inv(K + regularization_param * num_rows * I)
+    return W
diff --git a/gym_socks/kernel/probability.py b/gym_socks/kernel/probability.py
index 3bdb8ed..a56b1c5 100644
--- a/gym_socks/kernel/probability.py
+++ b/gym_socks/kernel/probability.py
@@ -3,9 +3,12 @@
 from functools import partial
 
 from gym_socks.kernel.metrics import regularized_inverse
-from gym_socks.kernel.metrics import rbf_kernel
 from gym_socks.kernel.metrics import delta_kernel
 
+from gym_socks.kernel.metrics import rbf_kernel
+
+from gym_socks.utils.validation import check_matrix
+
 
 def maximum_mean_discrepancy(
     X, Y, kernel_fn=None, biased: bool = False, squared: bool = False
@@ -36,8 +39,9 @@ def maximum_mean_discrepancy(
         X: The observations from distribution P organized in ROWS.
         Y: The observations from distribution Q organized in ROWS.
         kernel_fn: The kernel function is a function that returns an ndarray where
-            each element is the pairwise evaluation of a kernel function. See sklearn.
-            metrics.pairwise for more info. The default is the RBF kernel.
+            each element is the pairwise evaluation of a kernel function. See
+            :py:mod:`sklearn.metrics.pairwise` for more info. The default is the RBF
+            kernel.
         biased: Whether to use the biased MMD.
         squared: Whether or not the result is squared.
 
@@ -49,8 +53,7 @@ def maximum_mean_discrepancy(
     m = len(X)
     n = len(Y)
 
-    if kernel_fn is None:
-        kernel_fn = partial(rbf_kernel, sigma=1)
+    kernel_fn = rbf_kernel if kernel_fn is None else kernel_fn
 
     GX = kernel_fn(X)
     GY = kernel_fn(Y)
@@ -83,8 +86,7 @@ def witness_function(X, Y, t, kernel_fn=None):
     m = len(X)
     n = len(Y)
 
-    if kernel_fn is None:
-        kernel_fn = partial(rbf_kernel, sigma=1)
+    kernel_fn = rbf_kernel if kernel_fn is None else kernel_fn
 
     c1 = 1 / m
     c2 = 1 / n
@@ -244,30 +246,6 @@ def probability_matrix(states: np.ndarray, X: np.ndarray, Y: np.ndarray):
 
     """
 
-    if states.ndim == 1:
-        raise ValueError(
-            "Expected 2D array for states, got 1D array instead. \n"
-            "Reshape the data using array.reshape(-1, 1) "
-            "if the data has only a single dimension "
-            "or array.reshape(1, -1) if there is only a single sample."
-        )
-
-    if X.ndim == 1:
-        raise ValueError(
-            "Expected 2D array for X, got 1D array instead. \n"
-            "Reshape the data using array.reshape(-1, 1) "
-            "if the data has only a single dimension "
-            "or array.reshape(1, -1) if there is only a single sample."
-        )
-
-    if Y.ndim == 1:
-        raise ValueError(
-            "Expected 2D array for Y, got 1D array instead. \n"
-            "Reshape the data using array.reshape(-1, 1) "
-            "if the data has only a single dimension "
-            "or array.reshape(1, -1) if there is only a single sample."
-        )
-
     Gx = delta_kernel(states, X)
     Gy = delta_kernel(states, Y)
 
diff --git a/gym_socks/policies/policy.py b/gym_socks/policies/policy.py
index 525da33..a33067f 100644
--- a/gym_socks/policies/policy.py
+++ b/gym_socks/policies/policy.py
@@ -1,7 +1,7 @@
 """Control policies.
 
 Note:
-    Policies ccan be either time-invariant or time-varying, and can be either open- or
+    Policies can be either time-invariant or time-varying, and can be either open- or
     closed-loop. Thus, the arguments to the :py:meth:`__call__` method should allow for
     ``time`` and ``state`` to be specified (if needed), and should be optional kwargs::
 
@@ -24,7 +24,7 @@ class BasePolicy(ABC):
     Instead, define a new class that inherits from :py:class:`BasePolicy`.
 
     The :py:meth:`__call__` method is the main point of entry for the policy classes.
-    All subclasses must implement a py:meth:`__call__` method. This makes the class
+    All subclasses must implement a :py:meth:`__call__` method. This makes the class
     callable, so that policies can be evaluated as::
 
         action = policy(state)
diff --git a/gym_socks/sampling/__init__.py b/gym_socks/sampling/__init__.py
index 30bd737..dcd9f4b 100644
--- a/gym_socks/sampling/__init__.py
+++ b/gym_socks/sampling/__init__.py
@@ -29,23 +29,19 @@
 
 """
 
-from gym_socks.sampling.sample import sample_generator
+from gym_socks.sampling.sample import sample_fn
 
-from gym_socks.sampling.sample import sample
-
-from gym_socks.sampling.sample import default_sampler
-from gym_socks.sampling.sample import default_trajectory_sampler
+from gym_socks.sampling.sample import space_sampler
 from gym_socks.sampling.sample import grid_sampler
-from gym_socks.sampling.sample import random_sampler
 
-from gym_socks.sampling.sample import repeat
+from gym_socks.sampling.sample import transition_sampler
+from gym_socks.sampling.sample import trajectory_sampler
+
 
 __all__ = [
-    "sample_generator",
-    "sample",
-    "default_sampler",
-    "default_trajectory_sampler",
+    "sample_fn",
+    "space_sampler",
     "grid_sampler",
-    "random_sampler",
-    "repeat",
+    "transition_sampler",
+    "trajectory_sampler",
 ]
diff --git a/gym_socks/sampling/sample.py b/gym_socks/sampling/sample.py
index ffdba0b..42a4d7b 100644
--- a/gym_socks/sampling/sample.py
+++ b/gym_socks/sampling/sample.py
@@ -1,202 +1,239 @@
 """Sampling methods."""
 
 from inspect import isgeneratorfunction
-from functools import partial, wraps
+from functools import wraps
 from itertools import islice
+from collections.abc import Generator
 
 import gym
 
 from gym_socks.envs.dynamical_system import DynamicalSystem
 
+import numpy as np
 
-def sample_generator(fun):
-    """Sample generator decorator.
 
-    Converts a sample function into a generator function. Any function that returns a
-    single observation (as a tuple) can be converted into a sample generator.
+def _repeat(generator, num: int = 1):
+    """Repeat sampler.
+
+    Repeats the output of a sample generator ``num`` times.
 
     Args:
-        fun: Sample function that returns or yields an observation.
+        generator: The sample generator function.
+        num: The number of times to repeat a sample.
 
-    Returns:
-        A function that can be used to ``islice`` a sample from the sample generator.
+    Yields:
+        A repeated sample.
 
-    Example:
+    """
 
-        >>> from itertools import islice
-        >>> from gym_socks.envs.sample import sample_generator
-        >>> @sample_generator
-        ... def custom_sampler(env, policy, sample_space):
-        ...     env.state = sample_space.sample()
-        ...     action = policy(state=state)
-        ...     next_state, *_ = env.step(action)
-        ...     yield (env.state, action, next_state)
-        >>> S = list(islice(custom_sampler(), 100))
+    while True:
+        for item in generator:
+            for _ in range(num):
+                yield item
 
-    """
 
-    @wraps(fun)
-    def _wrapper(*args, **kwargs):
-        if isgeneratorfunction(fun):
-            while True:
-                yield from fun(*args, **kwargs)
-        else:
-            while True:
-                yield fun(*args, **kwargs)
+class _SampleGenerator(Generator):
+    """Sample generator object.
 
-    return _wrapper
+    This class wraps a generator created via :py:func:`sample_fn` and adds
+    useful helper functions for manipulating the sampling function.
 
+    Args:
+        generator: The generator function to wrap.
 
-@sample_generator
-def random_sampler(sample_space: gym.Space):
-    """Random sampler.
+    """
 
-    Returns a random sample taken from the space. See :py:obj:`gym.spaces.Box` for more
-    information on the distributions used for sampling.
+    def __init__(self, generator):
+        self._generator = generator
 
-    Args:
-        sample_space: The space to sample from.
+    def send(self, value):
+        return self._generator.send(value)
 
-    Yields:
-        A sample from the ``sample_space``.
+    def throw(self, type=None, value=None, traceback=None):
+        raise StopIteration
 
-    """
+    def __enter__(self):
+        return self
 
-    yield sample_space.sample()
+    def __exit__(self, type=None, value=None, traceback=None):
+        return False
 
+    def repeat(self, num: int = 1):
+        """Modify the generator to repeat the output ``num`` times.
 
-@sample_generator
-def grid_sampler(grid: list):
-    """Grid sampler.
+        See: :py:func:`_repeat` for more information.
 
-    Returns a sample arranged on a uniformly-spaced grid. Use
-    :py:func:`make_grid_from_ranges` or :py:func:`make_grid_from_space` from
-    :py:mod:`gym_socks.utils.grid` to generate a grid of points.
+        Args:
+            num: The number of times to repeat the output before continuing.
 
-    Args:
-        grid: The grid of points.
+        Returns:
+            An instance of the :py:class:`_SampleGenerator` object.
 
-    Yields:
-        A sample from the ``sample_space``.
+        """
 
-    """
+        self._generator = _repeat(self._generator, num)
+        return self
 
-    for item in grid:
-        yield item
+    def sample(self, size: int = 1):
+        """Generate a sample from the generator of a certain size."""
 
+        return list(islice(self._generator, size))
 
-@sample_generator
-def repeat(sampler, num: int):
-    """Repeat sampler.
 
-    Repeats the output of a sample generator ``num`` times.
+def sample_fn(fn):
+    """Sample function decorator.
+
+    Converts a sample function into a sample generator.
 
     Args:
-        sampler: The sample generator function.
-        num: The number of times to repeat a sample.
+        fn: Sample function that returns or yields an observation.
 
-    Yields:
-        A repeated sample.
+    Returns:
+        A function that can be used to ``islice`` a sample from the sample generator.
+
+    Example:
+
+        >>> from gym_socks.envs.sample import sample_fn
+        >>> @sample_fn
+        ... def custom_sampler(env, policy):
+        ...     state = env.reset()
+        ...     action = policy(state=state)
+        ...     observation, *_ = env.step(action)
+        ...     yield state, action, observation
+        >>> sampler = custom_sampler(env, policy)
+        >>> S = sampler.sample(size=100)
 
     """
 
-    for item in sampler:
-        for i in range(num):
-            yield item
+    # The inner wrapper protects in case the wrapped function is a generator.
+    if isgeneratorfunction(fn):
 
+        @wraps(fn)
+        def _inner_wrapper(*args, **kwargs):
+            while True:
+                yield from fn(*args, **kwargs)
 
-def default_sampler(
-    state_sampler=None,
-    action_sampler=None,
-    env: DynamicalSystem = None,
-):
-    """Default trajectory sampler.
+    else:
+
+        @wraps(fn)
+        def _inner_wrapper(*args, **kwargs):
+            while True:
+                yield fn(*args, **kwargs)
+
+    @wraps(fn)
+    def _wrapper(*args, **kwargs):
+        return _SampleGenerator(_inner_wrapper(*args, **kwargs))
+
+    return _wrapper
+
+
+@sample_fn
+def space_sampler(space: gym.spaces.Space):
+    """Randomly sample from a space.
 
     Args:
-        state_sampler: The state space sampler.
-        action_sampler: The action space sampler.
-        env: The system to sample from.
+        space: A :py:obj:`gym.space.Space` that implements a :py:func:`sample` function.
 
-    Returns:
-        A generator function that yields system observations as tuples.
+    Yields:
+        A random sample from the space.
 
-    """
+    Example:
 
-    @sample_generator
-    def _sample_generator():
-        state = next(state_sampler)
-        action = next(action_sampler)
+        >>> from gym.spaces import Box
+        >>> from gym_socks.envs.sample import space_sampler
+        >>> sampler = space_sampler(Box(low=-1, high=1, shape=(2,), dtype=float))
+        >>> S = sampler.sample(size=100)
 
-        env.state = state
-        next_state, *_ = env.step(action=action)
+    """
 
-        yield (state, action, next_state)
+    yield space.sample()
 
-    return _sample_generator
 
+@sample_fn
+def grid_sampler(grid_points: np.ndarray):
+    """Sample from a set of pre-defined grid points.
 
-def default_trajectory_sampler(
-    state_sampler=None,
-    action_sampler=None,
-    env: DynamicalSystem = None,
-    time_horizon: int = 1,
-):
-    """Default trajectory sampler.
+    Hint:
+        Use :py:func:`boxgrid` or :py:func:`cartesian` from
+        :py:mod:`gym_socks.utils.grid` to generate a grid of points.
 
     Args:
-        state_sampler: The state space sampler.
-        action_sampler: The action space sampler.
-        env: The system to sample from.
-        time_horizon: The time horizon to simulate over.
+        grid_points: A collection of grid points in a list or array.
 
-    Returns:
-        A generator function that yields system observations as tuples.
+    Yields:
+        A point in the grid. Note that points are yielded in the order they are given.
+
+    Example:
+
+        >>> import numpy as np
+        >>> from gym_socks.utils.grid import cartesian
+        >>> from gym_socks.envs.sample import grid_sampler
+        >>> grid = cartesian(np.linspace(-1, 1, 3), np.linspace(-2, 2, 3))
+        >>> sampler = grid_sampler(grid)
+        >>> S = sampler.sample(size=100)
 
     """
 
-    @sample_generator
-    def _sample_generator():
-        state = next(state_sampler)
+    for item in grid_points:
+        yield item
 
-        state_sequence = []
-        action_sequence = []
 
-        env.state = state
+@sample_fn
+def transition_sampler(
+    env: gym.Env,
+    state_sampler,
+    action_sampler,
+):
+    """Transition sampler."""
 
-        time = 0
-        for t in range(time_horizon):
-            action = next(action_sampler)
-            next_state, *_ = env.step(time=t, action=action)
+    initial_condition = next(state_sampler)
+    env.reset(initial_condition)
 
-            state_sequence.append(next_state)
-            action_sequence.append(action)
+    action = next(action_sampler)
+    env.step(action=action)
+    next_state = env.state
 
-            time += 1
+    yield initial_condition, action, next_state
 
-        yield (state, action_sequence, state_sequence)
 
-    return _sample_generator
+@sample_fn
+def observation_sampler(
+    env: gym.Env,
+    state_sampler,
+    action_sampler,
+):
+    """Observation sampler."""
 
+    initial_condition = next(state_sampler)
+    env.reset(initial_condition)
 
-def sample(sampler=None, sample_size: int = None, *args, **kwargs):
-    """Generate a sample using the sample generator.
+    action = next(action_sampler)
+    obs, *_ = env.step(action=action)
 
-    Args:
-        sampler: Sample generator function.
-        sample_size: Size of the sample.
+    yield initial_condition, action, obs
 
-    Returns:
-        list of tuples
 
-    """
+@sample_fn
+def trajectory_sampler(
+    env: gym.Env,
+    state_sampler,
+    action_sampler,
+    time_horizon: int = 1,
+):
+    """Trajectory sampler."""
 
-    if sampler is None:
-        raise ValueError("Must supply a sample function.")
+    initial_condition = next(state_sampler)
+    env.reset(initial_condition)
 
-    if sample_size is None:
-        raise ValueError("Must supply a sample size.")
+    state_sequence = []
+    action_sequence = []
+
+    for t in range(time_horizon):
+        action = next(action_sampler)
+        env.step(time=t, action=action)
+        next_state = env.state
 
-    if not isinstance(sample_size, int) or sample_size < 1:
-        raise ValueError("Sample size must be a positive integer.")
+        state_sequence.append(next_state)
+        action_sequence.append(action)
 
-    return list(islice(sampler(*args, **kwargs), sample_size))
+    yield initial_condition, np.array(action_sequence), np.array(state_sequence)
diff --git a/gym_socks/utils/__init__.py b/gym_socks/utils/__init__.py
index d856c76..5d2b120 100644
--- a/gym_socks/utils/__init__.py
+++ b/gym_socks/utils/__init__.py
@@ -1,4 +1,4 @@
-__all__ = ["logging", "normalize", "indicator_fn", "generate_batches"]
+__all__ = ["logging", "normalize", "indicator_fn", "batch_generator"]
 
 import gym
 
diff --git a/gym_socks/utils/batch.py b/gym_socks/utils/batch.py
index 7adcd26..9eeb03b 100644
--- a/gym_socks/utils/batch.py
+++ b/gym_socks/utils/batch.py
@@ -1,14 +1,16 @@
-def generate_batches(num_elements, batch_size):
+from collections.abc import Sequence
+
+
+def batch_generator(s: Sequence, size: int):
     """Generate batches.
 
-    Batch generation function to split a list into smaller batches (slices). Generates
-    `slice` objects, which can be iterated over in a for loop.
+    Batch generation function to split a sequence into smaller batches (slices).
+    Generates `slice` objects, which can be iterated over in a for loop.
 
     Args:
-        num_elements: Length of the list.
-        batch_size: Maximum size of the batches. If the last batch is smaller than
-            `batch_size`, then the final batch will be the length of the remaining
-            elements.
+        s: A sequence.
+        size: Size of the batches. If the last batch is smaller than `size`, then the
+            final batch will be the length of the remaining elements.
 
     Yields:
         A slice object.
@@ -17,12 +19,12 @@ def generate_batches(num_elements, batch_size):
 
     start = 0
 
-    for _ in range(num_elements // batch_size):
-        end = start + batch_size
+    for _ in range(len(s) // size):
+        end = start + size
 
         yield slice(start, end)
 
         start = end
 
-    if start < num_elements:
-        yield slice(start, num_elements)
+    if start < len(s):
+        yield slice(start, len(s))
diff --git a/gym_socks/utils/grid.py b/gym_socks/utils/grid.py
index d4a352a..98d6847 100644
--- a/gym_socks/utils/grid.py
+++ b/gym_socks/utils/grid.py
@@ -6,11 +6,11 @@
 from operator import mul
 
 
-def make_grid_from_ranges(xi: list) -> list:
-    """Create a grid of points from a list of ranges.
+def cartesian(*xi, out=None):
+    """Create a grid of points from ranges.
 
     Args:
-        xi: List of ranges.
+        xi: Ranges.
 
     Returns:
         Grid of points (the product of all points in ranges).
@@ -18,103 +18,91 @@ def make_grid_from_ranges(xi: list) -> list:
     Example:
 
         >>> import numpy as np
-        >>> from gym_socks.utils.grid import make_grid_from_ranges
-        >>> grid = make_grid_from_ranges([np.linspace(-1, 1, 5), np.linspace(-1, 1, 5)])
+        >>> from gym_socks.utils.grid import cartesian
+        >>> grid = cartesian(np.linspace(1, 3, 3), np.linspace(4, 6, 3))
+        array([[1.0, 4.0],
+               [1.0, 5.0],
+               [1.0, 6.0],
+               [2.0, 4.0],
+               [2.0, 5.0],
+               [2.0, 6.0],
+               [3.0, 4.0],
+               [3.0, 5.0],
+               [3.0, 6.0]])
 
     """
 
     len_xi = len(xi)
-
     dtype = np.result_type(*xi)
-    result = np.empty([len(a) for a in xi] + [len_xi], dtype=dtype)
+
+    if out is None:
+        out = np.empty([len(a) for a in xi] + [len_xi], dtype=dtype)
 
     for i, x in enumerate(np.ix_(*xi)):
-        result[..., i] = x
+        out[..., i] = x
+
+    out = out.reshape(-1, len_xi)
 
-    return result.reshape(-1, len_xi)
+    return out
 
 
-def make_grid_from_space(sample_space: gym.spaces.Box, resolution: int) -> list:
-    """Create a grid from a bounded sample space.
+def boxgrid(space: gym.spaces.Box, resolution: tuple, out=None):
+    """Returns a coarse grid from a bounded Box space.
+
+    Constructs a coarse grid from a bounded box in :math:`\mathbb{R}^n`. A
+    :py:class:`Box` is the Cartesian product of :math:`n` intervals, and we create a
+    coarse grid of those intervals using a specified :py:obj:`resolution`.
 
     Args:
-        sample_space: Bounded sample space.
-        resolution: The gid resolution, specified either as an integer for all
-            dimensions or as a list of integers, one per dimension.
+        space: The Box space to create a grid over. Must be bounded.
+        resolution: The grid resolution. Can either be a single integer (applies to all
+            dimensions) or a tuple of integers, which is of length :math:`n`.
 
     Returns:
         Grid of points (the product of all points in ranges).
 
     Example:
 
-        >>> from gym.spaces import Box
         >>> import numpy as np
-        >>> from gym_socks.utils.grid import make_grid_from_space
-        >>> sample_space = Box(low=-1, high=1, shape=(2,), dtype=np.float32)
-        >>> grid = make_grid_from_space(sample_space, [3, 2])
+        >>> from gym.spaces import Box
+        >>> from gym_socks.utils.grid import boxgrid
+        >>> space = Box(low=-1, high=1, shape=(2,), dtype=np.float32)
+        >>> boxgrid(space, (3, 5))
+        array([[-1. , -1. ],
+               [-1. , -0.5],
+               [-1. ,  0. ],
+               [-1. ,  0.5],
+               [-1. ,  1. ],
+               [ 0. , -1. ],
+               [ 0. , -0.5],
+               [ 0. ,  0. ],
+               [ 0. ,  0.5],
+               [ 0. ,  1. ],
+               [ 1. , -1. ],
+               [ 1. , -0.5],
+               [ 1. ,  0. ],
+               [ 1. ,  0.5],
+               [ 1. ,  1. ]])
 
     """
 
-    assert sample_space.is_bounded(), "space must be bounded."
-
-    low = sample_space.low
-    high = sample_space.high
+    assert space.is_bounded(), "space must be bounded."
 
     if np.isscalar(resolution):
-        xi = np.linspace(low, high, resolution, axis=-1)
-        grid = make_grid_from_ranges(xi)
+        xi = np.linspace(space.low, space.high, resolution, axis=-1, dtype=space.dtype)
+        out = cartesian(*xi, out=out)
 
     else:
         assert (
-            np.shape(resolution) == sample_space.shape
-        ), "resolution.shape doesn't match sample_space.shape"
+            np.shape(resolution) == space.shape
+        ), "resolution shape doesn't match space shape"
+
         xi = []
         for i, value in enumerate(resolution):
-            xi.append(np.linspace(low[i], high[i], value))
-
-        grid = make_grid_from_ranges(xi)
+            xi.append(
+                np.linspace(space.low[i], space.high[i], value, dtype=space.dtype)
+            )
 
-    return grid
-
-
-def grid_size_from_ranges(xi: list) -> int:
-    """Returns the size of a grid sample based on ranges.
-
-    Args:
-        xi: List of ranges.
-
-    Returns:
-        The product of the lengths of the ranges.
-
-    """
-
-    return reduce(mul, map(len, xi), 1)
-
-
-def grid_size_from_space(sample_space: gym.spaces.Box, resolution: int) -> int:
-    """Returns the size of a grid sample based on sample space and resolution.
-
-    Args:
-        sample_space: The sample space.
-        resolution: The gid resolution, specified either as an integer for all
-            dimensions or as a list of integers, one per dimension.
-
-    Returns:
-        The product of the lengths of the ranges.
-
-    """
-
-    assert sample_space.is_bounded(), "space must be bounded."
-
-    low = sample_space.low
-    high = sample_space.high
-
-    if np.isscalar(resolution):
-        resolution = np.full(shape=sample_space.shape, fill_value=resolution)
-
-    else:
-        assert (
-            np.shape(resolution) == sample_space.shape
-        ), "resolution.shape doesn't match sample_space.shape"
+        out = cartesian(*xi, out=out)
 
-    return reduce(mul, resolution, 1)
+    return out
diff --git a/gym_socks/utils/validation.py b/gym_socks/utils/validation.py
new file mode 100644
index 0000000..fccdb7f
--- /dev/null
+++ b/gym_socks/utils/validation.py
@@ -0,0 +1,98 @@
+import numpy as np
+
+
+def check_array(
+    array: np.ndarray,
+    dtype: np.dtype = np.float64,
+    order=None,
+    copy: bool = False,
+    ensure_2d: bool = True,
+    ensure_finite: bool = True,
+):
+    """Validate array.
+
+    Performs checks to ensure that the array is valid.
+
+    Note:
+
+        This function is intended as a simple replacement for the
+        :py:func:`sklearn.metrics.pairwise.check_array` function. Unlike
+        :py:mod:`sklearn`, this function does not check sparse input data, and does not
+        do sophisticated type checking or upcasting.
+
+    Args:
+        array: The array to be validated.
+        dtype: The data type of the resulting array.
+        order: The memory layout of the resulting array.
+        copy: Whether to create a forced copy of ``array``.
+        ensure_2d: Whether to raise an error if the array is not 2D.
+        ensure_finite: Whether to raise an error if the array is not finite.
+
+    Returns:
+        The validated array.
+
+    """
+
+    array_orig = array
+
+    # Convert array to numpy array.
+    array = np.asarray(array, dtype=dtype, order=order)
+
+    # Raise an error if the array is not 2D.
+    if ensure_2d is True:
+
+        if array.ndim == 0:
+            raise ValueError(
+                "Expected 2D array, got scalar array instead.\n"
+                "Reshape the data using array.reshape(-1, 1) "
+                "if the data has only a single dimension "
+                "or array.reshape(1, -1) if there is only a single sample."
+            )
+        elif array.ndim == 1:
+            raise ValueError(
+                "Expected 2D array, got 1D array instead.\n"
+                "Reshape the data using array.reshape(-1, 1) "
+                "if the data has only a single dimension "
+                "or array.reshape(1, -1) if there is only a single sample."
+            )
+
+    # Raise an error if the array elements are not finite.
+    if ensure_finite is True:
+        if array.dtype.kind in "fc":
+            if not np.isfinite(array).all():
+                raise ValueError("Input contains infinity or NaN.")
+
+    # Create a copy if requested.
+    if copy and np.may_share_memory(array, array_orig):
+        array = np.array(array, dtype=dtype, order=order)
+
+    return array
+
+
+def check_matrix(
+    A: np.ndarray,
+    ensure_square: bool = False,
+    copy: bool = False,
+):
+    """Validate a matrix.
+
+    Performs checks to ensure that a matrix is valid.
+
+    Args:
+        A: The matrix to validate.
+        ensure_square: Whether to raise an error if the matrix is not square.
+        copy: Whether to create a forced copy of ``A``.
+
+    Returns:
+        The validated matrix.
+
+    """
+
+    A = check_array(A, copy=copy)
+
+    if ensure_square is True:
+        m, n = np.shape(A)
+        if m != n:
+            raise ValueError("Matrix must be square.")
+
+    return A


Test patch:
diff --git a/gym_socks/algorithms/identification/tests/test_kernel_linear_id.py b/gym_socks/algorithms/identification/tests/test_kernel_linear_id.py
index 786cbcf..11c4f94 100644
--- a/gym_socks/algorithms/identification/tests/test_kernel_linear_id.py
+++ b/gym_socks/algorithms/identification/tests/test_kernel_linear_id.py
@@ -11,22 +11,18 @@
 from gym_socks.envs import CWH6DEnv
 from gym_socks.policies import RandomizedPolicy
 
-from gym_socks.sampling import sample
-from gym_socks.sampling import sample_generator
+from gym_socks.sampling import sample_fn
 
 
+@sample_fn
 def dummy_sampler(env, policy, sample_space):
-    @sample_generator
-    def _sample_generator():
-        state = sample_space.sample()
-        action = policy(time=0, state=state)
+    state = sample_space.sample()
+    action = policy(time=0, state=state)
 
-        env.state = state
-        next_state, *_ = env.step(action=action)
+    env.reset(state)
+    next_state, *_ = env.step(action=action)
 
-        yield (state, action, next_state)
-
-    return _sample_generator
+    yield state, action, next_state
 
 
 class TestKernelLinearIdentificationAlgorithm(unittest.TestCase):
@@ -36,14 +32,9 @@ def test_linear_identification_CWH4D(cls):
         env = CWH4DEnv()
         env.seed(seed=0)
 
-        S = sample(
-            sampler=dummy_sampler(
-                env=env,
-                policy=RandomizedPolicy(action_space=env.action_space),
-                sample_space=env.state_space,
-            ),
-            sample_size=1000,
-        )
+        S = dummy_sampler(
+            env, RandomizedPolicy(action_space=env.action_space), env.state_space
+        ).sample(size=1000)
 
         regularization_param = 1e-9
 
@@ -59,14 +50,9 @@ def test_linear_identification_CWH6D(cls):
         env = CWH6DEnv()
         env.seed(seed=0)
 
-        S = sample(
-            sampler=dummy_sampler(
-                env=env,
-                policy=RandomizedPolicy(action_space=env.action_space),
-                sample_space=env.state_space,
-            ),
-            sample_size=1000,
-        )
+        S = dummy_sampler(
+            env, RandomizedPolicy(action_space=env.action_space), env.state_space
+        ).sample(size=1000)
 
         regularization_param = 1e-9
 
diff --git a/gym_socks/algorithms/reach/tests/test_monte_carlo.py b/gym_socks/algorithms/reach/tests/test_monte_carlo.py
index 6362817..194c8f5 100644
--- a/gym_socks/algorithms/reach/tests/test_monte_carlo.py
+++ b/gym_socks/algorithms/reach/tests/test_monte_carlo.py
@@ -5,17 +5,16 @@
 import gym
 import gym_socks
 
-from scipy.constants.codata import unit
+from scipy.constants import unit
 
 from gym_socks.envs.integrator import NDIntegratorEnv
 from gym_socks.envs.dynamical_system import DynamicalSystem
 from gym_socks.policies import ZeroPolicy
 from gym_socks.policies import RandomizedPolicy
 
-from gym_socks.sampling import sample
-from gym_socks.sampling import default_trajectory_sampler
-from gym_socks.sampling import random_sampler
-from gym_socks.sampling import sample_generator
+from gym_socks.sampling import sample_fn
+from gym_socks.sampling import space_sampler
+from gym_socks.sampling import trajectory_sampler
 from gym_socks.sampling.transform import transpose_sample
 
 from gym_socks.algorithms.reach.common import _tht_step, _fht_step
@@ -44,7 +43,7 @@ def make_tube(
     return tube
 
 
-@sample_generator
+@sample_fn
 def zero_sampler(sample_space):
     yield np.zeros(shape=sample_space.shape, dtype=sample_space.dtype)
 
@@ -88,16 +87,21 @@ def test_trajectory_indicator_inside(cls, mock_generate_disturbance):
         )
 
         sample_size = 5
-
-        S = sample(
-            sampler=default_trajectory_sampler(
-                state_sampler=random_sampler(sample_space=sample_space),
-                action_sampler=zero_sampler(sample_space=env.action_space),
-                env=env,
-                time_horizon=cls.time_horizon,
-            ),
-            sample_size=sample_size,
-        )
+        state_sampler = space_sampler(space=sample_space)
+        action_sampler = zero_sampler(sample_space=env.action_space)
+        S = trajectory_sampler(
+            env, state_sampler, action_sampler, cls.time_horizon
+        ).sample(size=sample_size)
+
+        # S = sample(
+        #     sampler=default_trajectory_sampler(
+        #         state_sampler=random_sampler(sample_space=sample_space),
+        #         action_sampler=zero_sampler(sample_space=env.action_space),
+        #         env=env,
+        #         time_horizon=cls.time_horizon,
+        #     ),
+        #     sample_size=sample_size,
+        # )
 
         X, _, Y = transpose_sample(S)
         full_trajectories = [(X[i], *Y[i]) for i in range(sample_size)]
@@ -132,15 +136,21 @@ def test_trajectory_indicator_outside(cls, mock_generate_disturbance):
 
         sample_size = 5
 
-        S = sample(
-            sampler=default_trajectory_sampler(
-                state_sampler=random_sampler(sample_space=sample_space),
-                action_sampler=zero_sampler(sample_space=env.action_space),
-                env=env,
-                time_horizon=cls.time_horizon,
-            ),
-            sample_size=sample_size,
-        )
+        state_sampler = space_sampler(space=sample_space)
+        action_sampler = zero_sampler(sample_space=env.action_space)
+        S = trajectory_sampler(
+            env, state_sampler, action_sampler, cls.time_horizon
+        ).sample(size=sample_size)
+
+        # S = sample(
+        #     sampler=default_trajectory_sampler(
+        #         state_sampler=random_sampler(sample_space=sample_space),
+        #         action_sampler=zero_sampler(sample_space=env.action_space),
+        #         env=env,
+        #         time_horizon=cls.time_horizon,
+        #     ),
+        #     sample_size=sample_size,
+        # )
 
         X, _, Y = transpose_sample(S)
         full_trajectories = [(X[i], *Y[i]) for i in range(sample_size)]
diff --git a/gym_socks/envs/tests/test_cwh.py b/gym_socks/envs/tests/test_cwh.py
index a55043d..ad865e3 100644
--- a/gym_socks/envs/tests/test_cwh.py
+++ b/gym_socks/envs/tests/test_cwh.py
@@ -38,7 +38,7 @@ def test_known_trajectory(cls, mock_generate_disturbance):
 
         env = cls.env
 
-        env.state = np.array([-0.1, -0.1, 0, 0], dtype=np.float32)
+        env.reset([-0.1, -0.1, 0, 0])
         action = np.array([0, 0], dtype=np.float32)
 
         trajectory = []
@@ -109,7 +109,7 @@ def test_known_trajectory(cls, mock_generate_disturbance):
 
         env = cls.env
 
-        env.state = np.array([-0.1, -0.1, 0, 0, 0, 0], dtype=np.float32)
+        env.reset([-0.1, -0.1, 0, 0, 0, 0])
         action = np.array([0, 0, 0], dtype=np.float32)
 
         trajectory = []
diff --git a/gym_socks/envs/tests/test_integrator.py b/gym_socks/envs/tests/test_integrator.py
index 1e20feb..82f6fc6 100644
--- a/gym_socks/envs/tests/test_integrator.py
+++ b/gym_socks/envs/tests/test_integrator.py
@@ -5,6 +5,9 @@
 
 from gym_socks.envs.integrator import NDIntegratorEnv
 
+from gym_socks.policies import ZeroPolicy
+from gym_socks.envs.dynamical_system import simulate
+
 import numpy as np
 
 
@@ -21,17 +24,20 @@ def test_known_trajectory(cls, mock_generate_disturbance):
         env = cls.env
         env.sampling_time = 0.25
 
-        env.state = np.array([0.1, 0.1], dtype=np.float32)
-        action = np.array([0], dtype=np.float32)
+        # env.reset([0.1, 0.1])
+        # action = np.array([0], dtype=np.float32)
+
+        # trajectory = []
+        # trajectory.append(env.state)
 
-        trajectory = []
-        trajectory.append(env.state)
+        # for t in range(16):
+        #     obs, cost, done, _ = env.step(time=t, action=action)
+        #     trajectory.append(obs)
 
-        for t in range(16):
-            obs, cost, done, _ = env.step(time=t, action=action)
-            trajectory.append(obs)
+        # trajectory = np.array(trajectory)
 
-        trajectory = np.array(trajectory)
+        policy = ZeroPolicy(env.action_space)
+        trajectory, _ = simulate(env, 16, [0.1, 0.1], policy, True)
 
         groundTruth = np.array(
             [
@@ -66,17 +72,20 @@ def test_euler_approximation(cls, mock_generate_disturbance):
         env._euler = True
         env.sampling_time = 0.25
 
-        env.state = np.array([0.1, 0.1], dtype=np.float32)
-        action = np.array([0], dtype=np.float32)
+        # env.reset([0.1, 0.1])
+        # action = np.array([0], dtype=np.float32)
+
+        # trajectory = []
+        # trajectory.append(env.state)
 
-        trajectory = []
-        trajectory.append(env.state)
+        # for t in range(16):
+        #     obs, cost, done, _ = env.step(time=t, action=action)
+        #     trajectory.append(obs)
 
-        for t in range(16):
-            obs, cost, done, _ = env.step(time=t, action=action)
-            trajectory.append(obs)
+        # trajectory = np.array(trajectory)
 
-        trajectory = np.array(trajectory)
+        policy = ZeroPolicy(env.action_space)
+        trajectory, _ = simulate(env, 16, [0.1, 0.1], policy, True)
 
         groundTruth = np.array(
             [
diff --git a/gym_socks/envs/tests/test_world.py b/gym_socks/envs/tests/test_world.py
index 04cb22e..de8cfa5 100644
--- a/gym_socks/envs/tests/test_world.py
+++ b/gym_socks/envs/tests/test_world.py
@@ -54,5 +54,5 @@ def test_world(cls):
         for t in range(world.time_horizon):
             world.step()
 
-            action = policy(time=t, state=[env.state])
+            action = policy(time=t, state=env.state)
             env.step(time=t, action=action)
diff --git a/gym_socks/kernel/tests/test_kernel.py b/gym_socks/kernel/tests/test_kernel.py
index a98bc54..d8346eb 100644
--- a/gym_socks/kernel/tests/test_kernel.py
+++ b/gym_socks/kernel/tests/test_kernel.py
@@ -1,186 +1,155 @@
+from tokenize import group
 import unittest
 
 from functools import partial
 
-import gym
-
-from gym_socks.kernel.metrics import euclidean_distance
-from gym_socks.kernel.metrics import rbf_kernel as rbf_kernel_socks
-from gym_socks.kernel.metrics import regularized_inverse
-from gym_socks.kernel.metrics import delta_kernel
-from gym_socks.kernel.probability import probability_matrix
-
 import numpy as np
 
-from sklearn.metrics.pairwise import linear_kernel
-from sklearn.metrics.pairwise import polynomial_kernel
-from sklearn.metrics.pairwise import rbf_kernel
-from sklearn.metrics.pairwise import laplacian_kernel
-
-from sklearn.metrics.pairwise import euclidean_distances
-
-sklearn_kernel_list = [linear_kernel, polynomial_kernel, rbf_kernel, laplacian_kernel]
-
-
-class TestEuclideanDistance(unittest.TestCase):
-    def test_euclidean_distance(cls):
-        """Test that euclidean distance computes correctly."""
+import gym
 
-        X = np.arange(4).reshape((2, 2))
-        Y = np.arange(4).reshape((2, 2))
+from gym_socks.kernel.metrics import abel_kernel
+from gym_socks.kernel.metrics import delta_kernel
+from gym_socks.kernel.metrics import rbf_kernel
+from gym_socks.kernel.metrics import _hybrid_distances
+from gym_socks.kernel.metrics import hybrid_kernel
+from gym_socks.kernel.metrics import euclidean_distances
+from gym_socks.kernel.metrics import regularized_inverse
+from gym_socks.kernel.metrics import woodbury_inverse
 
-        groundTruth = np.array([[0.0, 8.0], [8.0, 0.0]], dtype=np.float32)
+from gym_socks.kernel.probability import probability_matrix
 
-        distance = euclidean_distance(Y, squared=True)
-        cls.assertTrue(np.allclose(distance, groundTruth))
 
-        distance = euclidean_distance(X, Y, squared=True)
-        cls.assertTrue(np.allclose(distance, groundTruth))
+class TestAbelKernel(unittest.TestCase):
+    def test_abel_kernel(cls):
+        """Test that Abel kernel computes correctly."""
 
-        groundTruth = euclidean_distances(X, Y, squared=True)
+        cls.assertTrue(True)
 
-        distance = euclidean_distance(Y, squared=True)
-        cls.assertTrue(np.allclose(distance, groundTruth))
 
-        distance = euclidean_distance(X, Y, squared=True)
-        cls.assertTrue(np.allclose(distance, groundTruth))
+class TestDeltaKernel(unittest.TestCase):
+    def test_delta_kernel(cls):
+        """Test that delta kernel computes correctly."""
 
+        X = np.array(
+            [
+                [0, 1],
+                [1, 1],
+                [1, 0],
+            ],
+            dtype=int,
+        )
 
-class TestRBFKernel(unittest.TestCase):
-    def test_rbf_kernel_same_size(cls):
-        """Test that RBF kernel computes correctly with same size inputs."""
+        G = delta_kernel(X)
 
-        Y = np.arange(4).reshape((2, 2))
+        cls.assertTrue(np.allclose(G, np.identity(3)))
 
-        groundTruth = np.array(
-            [[1.0, 0.77880078], [0.77880078, 1.0]],
-            dtype=np.float32,
+        Y = np.array(
+            [
+                [1, 1],
+                [2, 0],
+                [1, 0],
+                [1, 3],
+            ],
+            dtype=int,
         )
 
-        K = rbf_kernel_socks(Y)
-        cls.assertTrue(np.allclose(K, groundTruth))
+        G = delta_kernel(X, Y)
 
-        K = rbf_kernel_socks(Y, Y)
-        cls.assertTrue(np.allclose(K, groundTruth))
+        groundTruth = np.array([[0, 0, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0]], dtype=int)
 
-        K = rbf_kernel_socks(Y, distance_fn=euclidean_distances)
-        cls.assertTrue(np.allclose(K, groundTruth))
+        cls.assertTrue(np.allclose(G, groundTruth))
 
-        K = rbf_kernel_socks(Y, Y, distance_fn=euclidean_distances)
-        cls.assertTrue(np.allclose(K, groundTruth))
 
-    def test_rbf_kernel_different_size(cls):
-        """Test that RBF kernel computes correctly with different-sized inputs."""
+class TestHybridDistances(unittest.TestCase):
+    def test_hybrid_distances(cls):
+        """Test that hybrid distances computes correctly."""
 
-        X = np.arange(4).reshape((2, 2))
-        Y = np.arange(6).reshape((3, 2))
+        X = np.arange(12).reshape(6, 2)
+        Q = np.array([[0], [0], [0], [1], [2], [2]], dtype=int)
 
-        groundTruth = np.array(
-            [
-                [
-                    [1.0, 0.93941306, 0.77880078],
-                    [0.93941306, 1.0, 0.93941306],
-                    [0.77880078, 0.93941306, 1.0],
-                ],
-            ],
-            dtype=np.float32,
-        )
-
-        K = rbf_kernel_socks(Y)
-        cls.assertTrue(np.allclose(K, groundTruth))
-
-        K = rbf_kernel_socks(Y, Y)
-        cls.assertTrue(np.allclose(K, groundTruth))
+        D = _hybrid_distances(X, Q)
 
-        groundTruth = np.array(
-            [[1.0, 0.93941306, 0.77880078], [0.93941306, 1.0, 0.93941306]],
-            dtype=np.float32,
-        )
+        groundTruth = (2 / np.pi) * np.arctan(np.max(np.abs(X[0] - X[1])))
 
-        K = rbf_kernel_socks(X, Y)
-        cls.assertTrue(np.allclose(K, groundTruth))
+        cls.assertEqual(D[0, 1], groundTruth)
 
-        K = rbf_kernel_socks(X, Y, distance_fn=euclidean_distances)
-        cls.assertTrue(np.allclose(K, groundTruth))
+        groundTruth = (2 / np.pi) * np.arctan(np.max(np.abs(X[0] - X[2])))
 
-    def test_close_to_sklearn(cls):
-        """Assert that result is close to sklearn."""
+        cls.assertEqual(D[0, 2], groundTruth)
 
-        Y = np.arange(4).reshape((2, 2))
 
-        D = euclidean_distance(Y, squared=True)
+class TestHybridKernel(unittest.TestCase):
+    def test_hybrid_kernel(cls):
+        """Test that hybrid kernel computes correctly."""
 
-        groundTruth = rbf_kernel(Y, gamma=1 / (2 * np.median(D) ** 2))
+        X = np.arange(12).reshape(6, 2)
+        Q = np.array([[0], [0], [0], [1], [2], [2]], dtype=int)
 
-        K = rbf_kernel_socks(Y)
-        cls.assertTrue(np.allclose(K, groundTruth))
+        G = hybrid_kernel(X, Q)
 
-        K = rbf_kernel_socks(Y, Y)
-        cls.assertTrue(np.allclose(K, groundTruth))
+        cls.assertTrue(True)
 
 
 class TestRegularizedInverse(unittest.TestCase):
     def test_regularized_inverse(cls):
         """Test that regularized inverse computes correctly."""
 
-        Y = np.arange(4).reshape((2, 2))
+        X = np.arange(4).reshape((2, 2))
+
+        G = rbf_kernel(X, sigma=1)
 
         groundTruth = np.array(
             [[0.66676608, -0.0081415], [-0.0081415, 0.66676608]],
             dtype=np.float32,
         )
 
-        W = regularized_inverse(Y)
+        W = regularized_inverse(G, 1 / 4)
+
         cls.assertTrue(np.allclose(W, groundTruth))
 
-    def test_sklearn_kernels(cls):
-        """Sklearn kernels should work with regularized_inverse."""
 
-        Y = np.arange(4).reshape((2, 2))
+class TestWoodburyInverse(unittest.TestCase):
+    def test_woodbury_inverse(cls):
+        """Test that the Woodbury matrix identity computes correctly."""
 
-        for kernel_fn in sklearn_kernel_list:
-            with cls.subTest(msg=f"Testing with {type(kernel_fn)}."):
+        X = np.arange(12).reshape(6, 2)
+        groundTruth = regularized_inverse(X @ X.T, 1)
 
-                try:
+        A = 6 * np.identity(6)
+        C = np.identity(2)
+        W = woodbury_inverse(A, X, C, X.T)
 
-                    W = regularized_inverse(Y, Y, kernel_fn=kernel_fn)
+        cls.assertTrue(np.allclose(W, groundTruth))
 
-                except Exception as e:
-                    cls.fail(f"Kernel {type(kernel_fn)} raised an exception.")
+        X = np.arange(12).reshape(-1, 1)
+        groundTruth = regularized_inverse(X @ X.T, 1)
 
+        A = 12 * np.identity(12)
+        C = np.identity(1)
+        W = woodbury_inverse(A, X, C, X.T)
 
-class TestDeltaKernel(unittest.TestCase):
-    def test_delta_kernel(cls):
-        """Test that delta kernel computes correctly."""
+        cls.assertTrue(np.allclose(W, groundTruth))
 
-        X = np.array(
-            [
-                [0, 1],
-                [1, 1],
-                [1, 0],
-            ],
-            dtype=int,
-        )
+    def test_woodbury_inverse_precomputed(cls):
+        """Test that the Woodbury matrix identity computes correctly."""
 
-        G = delta_kernel(X)
+        X = np.arange(12).reshape(6, 2)
+        groundTruth = regularized_inverse(X @ X.T, 1)
 
-        cls.assertTrue(np.allclose(G, np.identity(3)))
+        A = (1 / 6) * np.identity(6)
+        C = np.identity(2)
+        W = woodbury_inverse(A, X, C, X.T, precomputed=True)
 
-        Y = np.array(
-            [
-                [1, 1],
-                [2, 0],
-                [1, 0],
-                [1, 3],
-            ],
-            dtype=int,
-        )
+        cls.assertTrue(np.allclose(W, groundTruth))
 
-        G = delta_kernel(X, Y)
+        X = np.arange(12).reshape(-1, 1)
+        groundTruth = regularized_inverse(X @ X.T, 1)
 
-        groundTruth = np.array([[0, 0, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0]], dtype=int)
+        A = (1 / 12) * np.identity(12)
+        C = np.identity(1)
+        W = woodbury_inverse(A, X, C, X.T, precomputed=True)
 
-        cls.assertTrue(np.allclose(G, groundTruth))
+        cls.assertTrue(np.allclose(W, groundTruth))
 
 
 class TestProbabilityMatrix(unittest.TestCase):
@@ -207,19 +176,11 @@ def test_probability_matrix_1D(cls):
 
         cls.assertTrue(np.allclose(P, groundTruth))
 
-    def test_probability_matrix_multi_dim(cls):
+    def test_probability_matrix_2D(cls):
         """Test that the probability matrix is correct for 2D discrete data."""
 
         # A list of all possible states.
-        states = np.array(
-            [
-                [0, 0],
-                [0, 1],
-                [1, 0],
-                [1, 1],
-            ],
-            dtype=int,
-        )
+        states = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=int)
 
         # Mock data. X represents the FROM state, Y represents the TO state.
         X = np.array([[0, 0], [0, 0], [0, 1], [1, 0], [1, 0], [1, 1]], dtype=int)
diff --git a/gym_socks/sampling/tests/test_sample.py b/gym_socks/sampling/tests/test_sample.py
index 815615c..9c113d2 100644
--- a/gym_socks/sampling/tests/test_sample.py
+++ b/gym_socks/sampling/tests/test_sample.py
@@ -7,19 +7,17 @@
 from gym_socks.policies import ConstantPolicy
 from gym_socks.policies import RandomizedPolicy
 
-from gym_socks.sampling.sample import sample_generator
-from gym_socks.sampling.sample import sample
-from gym_socks.sampling.sample import default_sampler
-from gym_socks.sampling.sample import default_trajectory_sampler
-from gym_socks.sampling.sample import random_sampler
+from gym_socks.sampling.sample import sample_fn
+from gym_socks.sampling.sample import space_sampler
 from gym_socks.sampling.sample import grid_sampler
-from gym_socks.sampling.sample import repeat
+from gym_socks.sampling.sample import transition_sampler
+from gym_socks.sampling.sample import trajectory_sampler
 
 from gym_socks.sampling.transform import transpose_sample
 from gym_socks.sampling.transform import flatten_sample
 
-from gym_socks.utils.grid import make_grid_from_ranges
-from gym_socks.utils.grid import make_grid_from_space
+from gym_socks.utils.grid import boxgrid
+from gym_socks.utils.grid import cartesian
 
 import numpy as np
 
@@ -38,17 +36,12 @@ def test_random_sampler(cls):
 
         sample_size = 10
 
-        state_sampler = random_sampler(sample_space=cls.env.state_space)
-        action_sampler = random_sampler(sample_space=cls.env.action_space)
+        state_sampler = space_sampler(space=cls.env.state_space)
+        action_sampler = space_sampler(space=cls.env.action_space)
 
-        S = sample(
-            sampler=default_sampler(
-                state_sampler=state_sampler,
-                action_sampler=action_sampler,
-                env=cls.env,
-            ),
-            sample_size=sample_size,
-        )
+        S = transition_sampler(
+            env=cls.env, state_sampler=state_sampler, action_sampler=action_sampler
+        ).sample(size=sample_size)
 
         X, U, Y = transpose_sample(S)
 
@@ -59,35 +52,24 @@ def test_random_sampler(cls):
     def test_grid_sampler(cls):
         """Grid sampler should output items on a grid of the correct size."""
 
-        # print(np.repeat(np.linspace(-1, 1, 5), 2, axis=0))
-
         sample_size = 27
 
-        state_sampler = repeat(
-            grid_sampler(
-                make_grid_from_space(
-                    sample_space=gym.spaces.Box(
-                        low=-1,
-                        high=1,
-                        shape=cls.env.state_space.shape,
-                        dtype=cls.env.state_space.dtype,
-                    ),
-                    resolution=3,
-                )
-            ),
-            num=3,
+        state_sample_space = gym.spaces.Box(
+            low=-1,
+            high=1,
+            shape=cls.env.state_space.shape,
+            dtype=cls.env.state_space.dtype,
         )
 
-        action_sampler = grid_sampler(make_grid_from_ranges([np.linspace(-1, 1, 3)]))
+        state_grid = boxgrid(space=state_sample_space, resolution=3)
+        action_grid = cartesian(np.linspace(-1, 1, 3))
 
-        S = sample(
-            sampler=default_sampler(
-                state_sampler=state_sampler,
-                action_sampler=action_sampler,
-                env=cls.env,
-            ),
-            sample_size=sample_size,
-        )
+        state_sampler = grid_sampler(grid_points=state_grid).repeat(num=3)
+        action_sampler = grid_sampler(grid_points=action_grid)
+
+        S = transition_sampler(
+            env=cls.env, state_sampler=state_sampler, action_sampler=action_sampler
+        ).sample(size=sample_size)
 
         X, U, Y = transpose_sample(S)
 
@@ -164,17 +146,12 @@ def test_default_sampler(cls):
 
         sample_size = 10
 
-        state_sampler = random_sampler(sample_space=cls.env.state_space)
-        action_sampler = random_sampler(sample_space=cls.env.action_space)
+        state_sampler = space_sampler(space=cls.env.state_space)
+        action_sampler = space_sampler(space=cls.env.action_space)
 
-        S = sample(
-            sampler=default_sampler(
-                state_sampler=state_sampler,
-                action_sampler=action_sampler,
-                env=cls.env,
-            ),
-            sample_size=sample_size,
-        )
+        S = transition_sampler(
+            env=cls.env, state_sampler=state_sampler, action_sampler=action_sampler
+        ).sample(size=sample_size)
 
         X, U, Y = transpose_sample(S)
 
@@ -188,18 +165,15 @@ def test_trajectory_sampler(cls):
         sample_size = 10
         time_horizon = 5
 
-        state_sampler = random_sampler(sample_space=cls.env.state_space)
-        action_sampler = random_sampler(sample_space=cls.env.action_space)
-
-        S = sample(
-            sampler=default_trajectory_sampler(
-                state_sampler=state_sampler,
-                action_sampler=action_sampler,
-                env=cls.env,
-                time_horizon=5,
-            ),
-            sample_size=sample_size,
-        )
+        state_sampler = space_sampler(space=cls.env.state_space)
+        action_sampler = space_sampler(space=cls.env.action_space)
+
+        S = trajectory_sampler(
+            env=cls.env,
+            state_sampler=state_sampler,
+            action_sampler=action_sampler,
+            time_horizon=5,
+        ).sample(size=sample_size)
 
         cls.assertEqual(np.array(S[0][0]).shape, (2,))
         cls.assertEqual(np.array(S[0][1]).shape, (5, 1))
@@ -222,22 +196,19 @@ def test_custom_sampler_as_function(cls):
 
         sample_size = 10
 
-        @sample_generator
+        @sample_fn
         def custom_sampler():
             state = cls.sample_space.sample()
             action = cls.env.action_space.sample()
 
-            cls.env.state = state
+            cls.env.reset(state)
             next_state, *_ = cls.env.step(action=action)
 
-            return (state, action, next_state)
+            return state, action, next_state
 
         #   ^^^^^^
 
-        S = sample(
-            sampler=custom_sampler,
-            sample_size=sample_size,
-        )
+        S = custom_sampler().sample(size=sample_size)
 
         X, U, Y = transpose_sample(S)
 
@@ -249,22 +220,19 @@ def test_custom_sampler_as_generator(cls):
 
         sample_size = 10
 
-        @sample_generator
+        @sample_fn
         def custom_sampler():
             state = cls.sample_space.sample()
             action = cls.env.action_space.sample()
 
-            cls.env.state = state
+            cls.env.reset(state)
             next_state, *_ = cls.env.step(action=action)
 
             yield (state, action, next_state)
 
         #   ^^^^^
 
-        S = sample(
-            sampler=custom_sampler,
-            sample_size=sample_size,
-        )
+        S = custom_sampler().sample(size=sample_size)
 
         X, U, Y = transpose_sample(S)
 
@@ -276,23 +244,20 @@ def test_custom_policy_sampler(cls):
 
         sample_size = 10
 
-        state_sampler = random_sampler(sample_space=cls.env.state_space)
+        state_sampler = space_sampler(space=cls.env.state_space)
         policy = ConstantPolicy(action_space=cls.env.action_space, constant=-1.0)
 
-        @sample_generator
+        @sample_fn
         def custom_sampler():
             state = next(state_sampler)
             action = policy()
 
-            cls.env.state = state
+            cls.env.reset(state)
             next_state, *_ = cls.env.step(action=action)
 
             yield (state, action, next_state)
 
-        S = sample(
-            sampler=custom_sampler,
-            sample_size=sample_size,
-        )
+        S = custom_sampler().sample(size=sample_size)
 
         X, U, Y = transpose_sample(S)
 
diff --git a/gym_socks/utils/tests/test_batch.py b/gym_socks/utils/tests/test_batch.py
index ec60f44..9237ef4 100644
--- a/gym_socks/utils/tests/test_batch.py
+++ b/gym_socks/utils/tests/test_batch.py
@@ -2,24 +2,24 @@
 
 import gym
 
-from gym_socks.utils.batch import generate_batches
+from gym_socks.utils.batch import batch_generator
 
 import numpy as np
 
 
 class TestGenerateBatches(unittest.TestCase):
-    """Test generate_batches."""
+    """Test batch_generator."""
 
-    def test_generate_batches(cls):
+    def test_batch_generator(cls):
         """Test generate batches."""
         idx = np.arange(10)
 
-        batches = generate_batches(num_elements=10, batch_size=2)
+        batches = batch_generator(s=idx, size=2)
         for i, batch in enumerate(batches):
             cls.assertTrue(len(idx[batch]) == 2)
             cls.assertTrue(np.all(np.equal(idx[batch], idx[i * 2 : i * 2 + 2])))
 
-        batches = generate_batches(num_elements=10, batch_size=5)
+        batches = batch_generator(s=idx, size=5)
         for i, batch in enumerate(batches):
             cls.assertTrue(len(idx[batch]) == 5)
             cls.assertTrue(np.all(np.equal(idx[batch], idx[i * 5 : i * 5 + 5])))
@@ -29,13 +29,13 @@ def test_batch_larger_than_set(cls):
 
         idx = np.arange(10)
 
-        batches = generate_batches(num_elements=14, batch_size=7)
+        batches = batch_generator(s=idx, size=7)
         batches.send(None)
         batch = batches.send(None)
         cls.assertTrue(len(idx[batch]) == 3)
         cls.assertTrue(np.all(np.equal(idx[batch], idx[-3:])))
 
-        batches = generate_batches(num_elements=10, batch_size=14)
+        batches = batch_generator(s=idx, size=14)
         batch = batches.send(None)
         cls.assertTrue(len(idx[batch]) == 10)
         cls.assertTrue(np.all(np.equal(idx[batch], idx[:])))
diff --git a/gym_socks/utils/tests/test_grid.py b/gym_socks/utils/tests/test_grid.py
index dbb2814..77de1fd 100644
--- a/gym_socks/utils/tests/test_grid.py
+++ b/gym_socks/utils/tests/test_grid.py
@@ -6,18 +6,14 @@
 
 import numpy as np
 
-from gym_socks.utils.grid import make_grid_from_ranges
-from gym_socks.utils.grid import make_grid_from_space
-from gym_socks.utils.grid import grid_size_from_ranges
-from gym_socks.utils.grid import grid_size_from_space
+from gym_socks.utils.grid import boxgrid
+from gym_socks.utils.grid import cartesian
 
 
 class TestGrid(unittest.TestCase):
-    def test_grid_from_ranges(cls):
-        """Should generate proper grid from ranges."""
-
-        xi = [np.linspace(-1, 1, 3), np.linspace(-1, 1, 3)]
-        result = make_grid_from_ranges(xi)
+    def test_cartesian(cls):
+        """Should generate all combinations of points (i.e. Cartesian product)."""
+        result = cartesian(np.linspace(-1, 1, 3), np.linspace(-1, 1, 3))
 
         groundTruth = [
             [-1.0, -1.0],
@@ -33,7 +29,7 @@ def test_grid_from_ranges(cls):
 
         cls.assertTrue(np.array_equiv(result, groundTruth))
 
-    def test_grid_from_space(cls):
+    def test_boxgrid(cls):
         """Should generate proper grid from space."""
 
         sample_space = gym.spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32)
@@ -50,26 +46,8 @@ def test_grid_from_space(cls):
             [1.0, 1.0],
         ]
 
-        result = make_grid_from_space(sample_space=sample_space, resolution=3)
+        result = boxgrid(space=sample_space, resolution=[3, 3])
         cls.assertTrue(np.array_equiv(result, groundTruth))
 
-        result = make_grid_from_space(sample_space=sample_space, resolution=[3, 3])
+        result = boxgrid(space=sample_space, resolution=3)
         cls.assertTrue(np.array_equiv(result, groundTruth))
-
-    def test_grid_size_from_space(cls):
-        """Should generate proper grid size from ranges."""
-
-        xi = [np.linspace(-1, 1, 3), np.linspace(-1, 1, 3)]
-        cls.assertEqual(grid_size_from_ranges(xi), 9)
-
-    def test_grid_size_from_space(cls):
-        """Should generate proper grid size from space."""
-
-        sample_space = gym.spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32)
-
-        cls.assertEqual(
-            grid_size_from_space(sample_space=sample_space, resolution=3), 9
-        )
-        cls.assertEqual(
-            grid_size_from_space(sample_space=sample_space, resolution=[3, 3]), 9
-        )