Fixes nightly builds moving to python 3.9.

- Removed tensorflow-datasets, it triggers trying to install etils[array_types] which doesn't exist. I am not sure when it existed. I tried 1.4.1 and 1.5.0 (latest).
- ubuntu_ci has python3.9 and 3.10.
- Did a successful build and colab test with TF 2.14.0, Reverb 2.13.0, and tfp-nightly (need because the last release is not compatibility with TF 2.14.0.
- Had to turn off `bandits.agents.linear_bandit_agent_test` test due to assertion error. Sees 2 somethings in the event log and expects 1. Given our issues turning this off is prudent. P1 bug opened on that team.

PiperOrigin-RevId: 569588372
Change-Id: Id403aa14226d586a72628120b4a59d9a9759da0a

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
+examples.cql_sac.kumar20.cql_sac_train_eval_test  # b/199583100
+bandits.agents.linear_bandit_agent_test  # b/302734159

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
+networks.actor_distribution_rnn_network_test  # b/237573967
+agents.ddpg.critic_rnn_network_test  # b/237573967
+bandits.agents.linear_bandit_agent_test  # b/302734159

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
 

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