diff --git a/README.md b/README.md
index 187df0e..db8fb96 100644
--- a/README.md
+++ b/README.md
@@ -2,7 +2,9 @@
 This is the code used for the paper [Towards Learning Universal Hyperparameter Optimizers with Transformers (NeurIPS 2022)](https://arxiv.org/abs/2205.13320).
 
 # Installation
-All base dependencies can be installed from `requirements.txt`. Afterwards, [T5X](https://github.com/google-research/t5x) must be manually installed.
+To get started, manually install [T5X](https://github.com/google-research/t5x). Afterwards, install OptFormer with ``pip install -e .``.
+
+
 
 # Usage
 
diff --git a/optformer/t5x/policies.py b/optformer/t5x/policies.py
index 6538876..b13e76a 100644
--- a/optformer/t5x/policies.py
+++ b/optformer/t5x/policies.py
@@ -387,9 +387,8 @@ def _thompson_sampling(
     return max_funs
 
   def update(
-      self, completed: vza.CompletedTrials, all_active: vza.ActiveTrials
+      self, completed: vza.CompletedTrials
   ) -> None:
-    del all_active
     completed_trials = []
     for trial in completed.trials:
       # A completed trial either has a final_measurement or is marked as
diff --git a/requirements.txt b/requirements.txt
index ca7c8c8..4effe59 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,12 +1,16 @@
 # T5X prerequisite.
-seqio  # Always compatible with most up-to-date version.
+t5
+seqio==0.0.13  # Always compatible with most up-to-date version.
 
 # OSS Vizier
-google-vizier # Use latest version.
+google-vizier==0.0.20 # Use latest version.
 
 # Proto
 protobuf>=3.6,<4.0
 
+# For the notebook
+matplotlib
+
 # Autodiff and numerical packages.
 numpy>=1.21.5
 tensorflow>=2.9.1