keep using tf-keras rather than keras-3 which is not compatible.

Also, moving tensorflow from [reverb] requirement to main requirments.

PiperOrigin-RevId: 592609539
Change-Id: Ife593ab7960dddbd545d33eb9a49dce60e27f1f6

README.md

@@ -108,6 +108,8 @@ documentation for the release is on
 ```shell
 $ pip install --user tf-agents[reverb]
 
+# Use keras-2
+$ export TF_USE_LEGACY_KERAS=1
 # Use this tag get the matching examples and colabs.
 $ git clone https://github.com/tensorflow/agents.git
 $ cd agents
@@ -120,6 +122,7 @@ by the pip dependency check, use the following pattern below at your own risk.
 
 ```shell
 $ pip install --user tensorflow
+$ pip install --user tf-keras
 $ pip install --user dm-reverb
 $ pip install --user tf-agents
 ```
@@ -143,8 +146,12 @@ tested against.
 To install the nightly build version, run the following:
 
 ```shell
+# Use keras-2
+$ export TF_USE_LEGACY_KERAS=1
+
 # `--force-reinstall helps guarantee the right versions.
 $ pip install --user --force-reinstall tf-nightly
+$ pip install --user --force-reinstall tf-keras-nightly
 $ pip install --user --force-reinstall tfp-nightly
 $ pip install --user --force-reinstall dm-reverb-nightly
 

docs/overview.md

@@ -29,6 +29,8 @@ documentation for the release is on
 ```shell
 $ pip install --user tf-agents[reverb]
 
+# Use keras-2
+$ export TF_USE_LEGACY_KERAS=1
 # Use this tag get the matching examples and colabs.
 $ git clone https://github.com/tensorflow/agents.git
 $ cd agents
@@ -41,6 +43,7 @@ by the pip dependency check, use the following pattern below at your own risk.
 
 ```shell
 $ pip install --user tensorflow
+$ pip install --user tf-keras
 $ pip install --user dm-reverb
 $ pip install --user tf-agents
 ```
@@ -64,8 +67,11 @@ tested against.
 To install the nightly build version, run the following:
 
 ```shell
+# Use keras-2
+$ export TF_USE_LEGACY_KERAS=1
 # `--force-reinstall helps guarantee the right versions.
 $ pip install --user --force-reinstall tf-nightly
+$ pip install --user --force-reinstall tf-keras-nightly
 $ pip install --user --force-reinstall tfp-nightly
 $ pip install --user --force-reinstall dm-reverb-nightly
 

docs/tutorials/10_checkpointer_policysaver_tutorial.ipynb

@@ -111,6 +111,19 @@
         "!pip install tf-agents[reverb]"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "WPuD0bMEY9Iz"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "# Keep using keras-2 (tf-keras) rather than keras-3 (keras).\n",
+        "os.environ['TF_USE_LEGACY_KERAS'] = '1'"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,

docs/tutorials/1_dqn_tutorial.ipynb

@@ -120,6 +120,19 @@
         "!pip install pyglet"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "UX0aSKBCYmj2"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "# Keep using keras-2 (tf-keras) rather than keras-3 (keras).\n",
+        "os.environ['TF_USE_LEGACY_KERAS'] = '1'"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -222,7 +235,7 @@
         "\n",
         "In Reinforcement Learning (RL), an environment represents the task or problem to be solved. Standard environments can be created in TF-Agents using `tf_agents.environments` suites. TF-Agents has suites for loading environments from sources such as the OpenAI Gym, Atari, and DM Control.\n",
         "\n",
-        "Load the CartPole environment from the OpenAI Gym suite. "
+        "Load the CartPole environment from the OpenAI Gym suite."
       ]
     },
     {
@@ -323,9 +336,9 @@
       "source": [
         "In the Cartpole environment:\n",
         "\n",
-        "-   `observation` is an array of 4 floats: \n",
+        "-   `observation` is an array of 4 floats:\n",
         "    -   the position and velocity of the cart\n",
-        "    -   the angular position and velocity of the pole \n",
+        "    -   the angular position and velocity of the pole\n",
         "-   `reward` is a scalar float value\n",
         "-   `action` is a scalar integer with only two possible values:\n",
         "    -   `0` — \"move left\"\n",
@@ -357,7 +370,7 @@
         "id": "4JSc9GviWUBK"
       },
       "source": [
-        "Usually two environments are instantiated: one for training and one for evaluation. "
+        "Usually two environments are instantiated: one for training and one for evaluation."
       ]
     },
     {
@@ -500,7 +513,7 @@
         "-   The desired outcome is keeping the pole balanced upright over the cart.\n",
         "-   The policy returns an action (left or right) for each `time_step` observation.\n",
         "\n",
-        "Agents contain two policies: \n",
+        "Agents contain two policies:\n",
         "\n",
         "-   `agent.policy` — The main policy that is used for evaluation and deployment.\n",
         "-   `agent.collect_policy` — A second policy that is used for data collection.\n"
@@ -834,7 +847,7 @@
       "source": [
         "# For the curious:\n",
         "# Uncomment to see what the dataset iterator is feeding to the agent.\n",
-        "# Compare this representation of replay data \n",
+        "# Compare this representation of replay data\n",
         "# to the collection of individual trajectories shown earlier.\n",
         "\n",
         "# iterator.next()"
@@ -967,7 +980,7 @@
         "id": "9pGfGxSH32gn"
       },
       "source": [
-        "Charts are nice. But more exciting is seeing an agent actually performing a task in an environment. \n",
+        "Charts are nice. But more exciting is seeing an agent actually performing a task in an environment.\n",
         "\n",
         "First, create a function to embed videos in the notebook."
       ]
@@ -1048,7 +1061,6 @@
   ],
   "metadata": {
     "colab": {
-      "collapsed_sections": [],
       "name": "DQN Tutorial.ipynb",
       "private_outputs": true,
       "provenance": [],

docs/tutorials/2_environments_tutorial.ipynb

@@ -112,6 +112,19 @@
         "!pip install tf-agents[reverb]\n"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "WPuD0bMEY9Iz"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "# Keep using keras-2 (tf-keras) rather than keras-3 (keras).\n",
+        "os.environ['TF_USE_LEGACY_KERAS'] = '1'"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -489,7 +502,7 @@
         "The interface for TF environments is defined in `environments/tf_environment.TFEnvironment` and looks very similar to the Python environments. TF Environments differ from Python envs in a couple of ways:\n",
         "\n",
         "* They generate tensor objects instead of arrays\n",
-        "* TF environments add a batch dimension to the tensors generated when compared to the specs. \n",
+        "* TF environments add a batch dimension to the tensors generated when compared to the specs.\n",
         "\n",
         "Converting the Python environments into TFEnvs allows tensorflow to parallelize operations. For example, one could define a `collect_experience_op` that collects data from the environment and adds to a `replay_buffer`, and a `train_op` that reads from the `replay_buffer` and trains the agent, and run them in parallel naturally in TensorFlow."
       ]
@@ -702,7 +715,6 @@
   ],
   "metadata": {
     "colab": {
-      "collapsed_sections": [],
       "name": "TF-Agents Environments Tutorial.ipynb",
       "private_outputs": true,
       "provenance": [],

docs/tutorials/3_policies_tutorial.ipynb

@@ -114,6 +114,19 @@
         "!pip install tf-agents"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "h3dkwi09ZQeJ"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "# Keep using keras-2 (tf-keras) rather than keras-3 (keras).\n",
+        "os.environ['TF_USE_LEGACY_KERAS'] = '1'"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -295,9 +308,9 @@
       "outputs": [],
       "source": [
         "action_spec = array_spec.BoundedArraySpec((2,), np.int32, -10, 10)\n",
-        "action_script = [(1, np.array([5, 2], dtype=np.int32)), \n",
+        "action_script = [(1, np.array([5, 2], dtype=np.int32)),\n",
         "                 (0, np.array([0, 0], dtype=np.int32)), # Setting `num_repeats` to 0 will skip this action.\n",
-        "                 (2, np.array([1, 2], dtype=np.int32)), \n",
+        "                 (2, np.array([1, 2], dtype=np.int32)),\n",
         "                 (1, np.array([3, 4], dtype=np.int32))]\n",
         "\n",
         "my_scripted_py_policy = scripted_py_policy.ScriptedPyPolicy(\n",
@@ -555,7 +568,7 @@
         "id": "BzoNGJnlibtz"
       },
       "source": [
-        "Note that in the above, actions are clipped to the range of the given action spec [-1, 1]. This is because a constructor argument of ActorPolicy clip=True by default. Setting this to false will return unclipped actions produced by the network. "
+        "Note that in the above, actions are clipped to the range of the given action spec [-1, 1]. This is because a constructor argument of ActorPolicy clip=True by default. Setting this to false will return unclipped actions produced by the network."
       ]
     },
     {
@@ -655,7 +668,7 @@
         "id": "OfaUrqRAoigk"
       },
       "source": [
-        "A policy wrapper can be used to wrap and modify a given policy, e.g. add noise. Policy wrappers are a subclass of Policy (Python/TensorFlow) and can therefore be used just like any other policy. "
+        "A policy wrapper can be used to wrap and modify a given policy, e.g. add noise. Policy wrappers are a subclass of Policy (Python/TensorFlow) and can therefore be used just like any other policy."
       ]
     },
     {
@@ -692,7 +705,6 @@
   ],
   "metadata": {
     "colab": {
-      "collapsed_sections": [],
       "name": "TF-Agents Policies Tutorial.ipynb",
       "private_outputs": true,
       "provenance": [],

docs/tutorials/4_drivers_tutorial.ipynb

@@ -105,6 +105,19 @@
         "!pip install tf-agents\n"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "WPuD0bMEY9Iz"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "# Keep using keras-2 (tf-keras) rather than keras-3 (keras).\n",
+        "os.environ['TF_USE_LEGACY_KERAS'] = '1'"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,

docs/tutorials/5_replay_buffers_tutorial.ipynb

@@ -103,6 +103,19 @@
         "!pip install tf-agents\n"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "WPuD0bMEY9Iz"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "# Keep using keras-2 (tf-keras) rather than keras-3 (keras).\n",
+        "os.environ['TF_USE_LEGACY_KERAS'] = '1'"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,

docs/tutorials/6_reinforce_tutorial.ipynb

@@ -117,6 +117,19 @@
         "!pip install pyglet xvfbwrapper\n"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "WPuD0bMEY9Iz"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "# Keep using keras-2 (tf-keras) rather than keras-3 (keras).\n",
+        "os.environ['TF_USE_LEGACY_KERAS'] = '1'"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -571,7 +584,7 @@
         "  # Use data from the buffer and update the agent's network.\n",
         "  iterator = iter(replay_buffer.as_dataset(sample_batch_size=1))\n",
         "  trajectories, _ = next(iterator)\n",
-        "  train_loss = tf_agent.train(experience=trajectories)  \n",
+        "  train_loss = tf_agent.train(experience=trajectories)\n",
         "\n",
         "  replay_buffer.clear()\n",
         "\n",
@@ -697,7 +710,6 @@
   "metadata": {
     "accelerator": "GPU",
     "colab": {
-      "collapsed_sections": [],
       "name": "6_reinforce_tutorial.ipynb",
       "provenance": [],
       "toc_visible": true

docs/tutorials/7_SAC_minitaur_tutorial.ipynb

@@ -112,6 +112,19 @@
         "!pip install pybullet"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "WPuD0bMEY9Iz"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "# Keep using keras-2 (tf-keras) rather than keras-3 (keras).\n",
+        "os.environ['TF_USE_LEGACY_KERAS'] = '1'"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {

docs/tutorials/8_networks_tutorial.ipynb

@@ -117,6 +117,19 @@
         "!pip install tf-agents"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "WPuD0bMEY9Iz"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "# Keep using keras-2 (tf-keras) rather than keras-3 (keras).\n",
+        "os.environ['TF_USE_LEGACY_KERAS'] = '1'"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,

docs/tutorials/9_c51_tutorial.ipynb

@@ -117,6 +117,19 @@
         "!pip install pyglet"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "WPuD0bMEY9Iz"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "# Keep using keras-2 (tf-keras) rather than keras-3 (keras).\n",
+        "os.environ['TF_USE_LEGACY_KERAS'] = '1'"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,