diff --git a/broken_tests.txt b/broken_tests.txt
index e558b2a36..76ebbace0 100644
--- a/broken_tests.txt
+++ b/broken_tests.txt
@@ -2,4 +2,5 @@ agents.cql.cql_sac_agent_test  # b/239448722
 environments.suite_atari_test # b/201537188
 examples.cql_sac.kumar20.dataset.file_utils_test  # b/199583100
 examples.cql_sac.kumar20.dataset.dataset_utils_test  # b/199583100
-examples.cql_sac.kumar20.cql_sac_train_eval_test  # b/199583100
\ No newline at end of file
+examples.cql_sac.kumar20.cql_sac_train_eval_test  # b/199583100
+bandits.agents.linear_bandit_agent_test  # b/302734159
\ No newline at end of file
diff --git a/broken_tests_gpu.txt b/broken_tests_gpu.txt
index fd4d6873c..594fd82e8 100644
--- a/broken_tests_gpu.txt
+++ b/broken_tests_gpu.txt
@@ -6,5 +6,6 @@ examples.cql_sac.kumar20.cql_sac_train_eval_test  # b/199583100
 replay_buffers.episodic_replay_buffer_driver_test  # b/199587793
 replay_buffers.episodic_replay_buffer_test  # b/199587793
 networks.q_rnn_network_test # b/237573967
-networks.actor_distribution_rnn_network_test # b/237573967
-agents.ddpg.critic_rnn_network_test # b/237573967
\ No newline at end of file
+networks.actor_distribution_rnn_network_test  # b/237573967
+agents.ddpg.critic_rnn_network_test  # b/237573967
+bandits.agents.linear_bandit_agent_test  # b/302734159
\ No newline at end of file
diff --git a/setup.py b/setup.py
index f245fe978..8357d6443 100644
--- a/setup.py
+++ b/setup.py
@@ -212,7 +212,6 @@ def get_test_packages():
       'opencv-python >= 3.4.1.15',
       'pybullet',
       'scipy >= 1.1.0',
-      'tensorflow_datasets',
   ]
   return test_packages
 
diff --git a/tools/docker/ubuntu_ci b/tools/docker/ubuntu_ci
index 415476558..9356e179c 100644
--- a/tools/docker/ubuntu_ci
+++ b/tools/docker/ubuntu_ci
@@ -12,7 +12,7 @@ FROM $base_image
 LABEL maintainer="tobyboyd@google.com"
 
 # Supports setting up a single version of python.
-ARG python_version="python3"
+ARG python_version="python3.9"
 ARG APT_COMMAND="apt-get -o Acquire::Retries=3 -y"
 
 # Stops tzdata from asking about timezones and blocking install on user input.
@@ -42,8 +42,11 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
       xvfb \
       freeglut3-dev \
       ffmpeg \
-      python3-opencv \
-      python3-dev \
+      python3.9-dev \
+      python3.10-dev \
+      # python >= 3.8 needs distutils for packaging.
+      python3.9-distutils \
+      python3.10-distutils \
       gfortran \
       libopenblas-dev \
       liblapack-dev