google-deepmind · copybara-service · Jun 29, 2026
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -28,9 +28,6 @@ on:
 jobs:
   build:
     runs-on: ubuntu-latest
-    env:
-      TEST_TMPDIR: '/tmp'
-      USE_BAZEL_VERSION: 8.2.1
     strategy:
       matrix:
         python-version: ["3.11", "3.12", "3.13"]
@@ -42,35 +39,16 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
-      - name: Install dependencies
+      - name: Install system dependencies
         run: |
-          set -x
-          pip install --upgrade pip setuptools wheel
-
-          # Build and install Riegeli from source
-          git clone --depth=1 https://github.com/google/riegeli.git /tmp/riegeli
-          cd /tmp/riegeli
-          # Downgrade protobuf in MODULE.bazel to avoid python runtime incompatibility with TFDS (which needs <7)
-          sed -i 's|version = "34.1"|version = "28.3"|' MODULE.bazel
-          if [ -f configure ]; then ./configure; fi
-          # Disable TensorFlow support in Riegeli python build to avoid local_config_tf errors
-          rm -rf python/riegeli/tensorflow
-          sed -i '/riegeli_dataset_ops/d' python/BUILD
-          sed -i 's|cd bazel-bin/python/build_pip_package.runfiles/com_google_riegeli/python|if [ -d bazel-bin/python/build_pip_package.runfiles/_main/python ]; then cd bazel-bin/python/build_pip_package.runfiles/_main/python; else cd bazel-bin/python/build_pip_package.runfiles/com_google_riegeli/python; fi|' python/build_pip_package.sh
-          grep "_main" python/build_pip_package.sh || { echo "Patch failed!"; exit 1; }
-          bazel build -c opt --@rules_python//python/config_settings:python_version=${{ matrix.python-version }} python:build_pip_package
-          ./bazel-bin/python/build_pip_package --dest /tmp/riegeli_dist --bdist || { ls -R bazel-bin/python/build_pip_package.runfiles; exit 1; }
-          pip install /tmp/riegeli_dist/*.whl
+          sudo apt-get update
+          sudo apt-get install -y libgmp-dev
 
-          # Go back to project root and install envlogger
-          cd $GITHUB_WORKSPACE
-          pip install .[tfds]
-          pip list
-          pip install "protobuf<7"
-          pip list
+      - name: Install Python dependencies
+        run: |
+          pip install --upgrade pip setuptools wheel
+          pip install -r envlogger/requirements.txt
 
       - name: Run tests
         run: |
-          # Find all Python test files (excluding Bazel-only cross-language tests),
-          # print their names and execute them in parallel with a maximum of 20 processes.
-          find envlogger -type f -name "*_test.py" ! -path "*/cross_language_test/*" -print0 | xargs -t -0 -n1 -P 20 python3
+          bazel test -c fastbuild --test_output=errors envlogger/...
diff --git a/README.md b/README.md
@@ -46,7 +46,6 @@ with envlogger.EnvLogger(
 Full example of random agent in Catch is available here:
 [random_agent_catch.py](https://github.com/google-deepmind/envlogger/tree/main/envlogger/examples/random_agent_catch.py)
 
-
 ### Step metadata
 
 Envlogger also allows to record custom metadata per step by defining a function
@@ -102,9 +101,10 @@ episode.
 Envlogger supports writing data that is directly compatible with
 [TFDS](http://www.tensorflow.org/datasets) and [RLDS].
 
-For that, you need to indicate the specs of your environment in terms of [TFDS
-Features](https://www.tensorflow.org/datasets/features) using an RLDS [DatasetConfig] like in the example
-(note that the config can include `step_metadata_info`and `episode metadata_info`).
+For that, you need to indicate the specs of your environment in terms of
+[TFDS Features](https://www.tensorflow.org/datasets/features) using an
+RLDS [DatasetConfig] like in the example (note that the config can include
+`step_metadata_info` and `episode metadata_info`).
 
 ```
   dataset_config = tfds.rlds.rlds_base.DatasetConfig(
@@ -132,7 +132,6 @@ envlogger.EnvLogger(
 
 A full example is available here [random_agent_catch.py](https://github.com/google-deepmind/envlogger/tree/main/envlogger/examples/tfds_random_agent_catch.py)
 
-
 [RLDS]: http://github.com/google-research/rlds
 [DatasetConfig]: https://github.com/tensorflow/datasets/blob/fdad1d9e8f1cb34389a336132b2f842cbc7aca57/tensorflow_datasets/rlds/rlds_base.py#L29
 [TFDS Backend]:https://github.com/deepmind/envlogger/blob/main/envlogger/backends/tfds_backend_writer.py
@@ -194,8 +193,7 @@ pip install envlogger[tfds]
 ##### Compiling from source
 
 For this option you will need to [install Bazel](https://docs.bazel.build/versions/main/install.html) and [GMP](https://gmplib.org/) (`libgmp-dev` on Debian-based systems).
-Please note that Bazel versions >4.0 are not supported. Our recommended version
-is [3.7.2](https://github.com/bazelbuild/bazel/releases/tag/3.7.2). Then:
+We recommend using [bazelisk](https://github.com/bazelbuild/bazelisk) to manage Bazel versions, which will automatically use the pinned version in `.bazelversion` (currently 7.1.0). Then:
 
 ```
 git clone https://github.com/deepmind/envlogger/
@@ -216,8 +214,10 @@ bazel test --test_output=errors envlogger/...
 
 ## Benchmarks
 
-Wrapping your environment with EnvLogger adds an approximately 2 millisecond overhead per environment step.
-See the following difference in distributions in the case of random agent on a 100 steps per second Catch environment (measured in milliseconds).
+Wrapping your environment with EnvLogger adds an approximately 2 millisecond
+overhead per environment step.
+See the following difference in distributions in the case of random agent on a
+100 steps per second Catch environment (measured in milliseconds).
 
 Percentiles      | 50th      | 95th       | 99th      | 99.9th     | mean (± std)
 ---------------- | --------- | -----------| ----------| -----------| -----------
@@ -226,7 +226,6 @@ w/ EnvLogger     | 12.65     | 14.39      | 15.94     | 19.43      | 12.88 (± 0
 
 <img src="docs/images/timings.png" width="40%">
 
-
 ## Acknowledgements
 
 We greatly appreciate all the support from the

diff --git a/WORKSPACE.bazel b/WORKSPACE.bazel
@@ -84,15 +84,13 @@ load("@rules_python//python:repositories.bzl", "py_repositories", "python_regist
 
 py_repositories()
 
-python_register_toolchains(
-    name = "python_3_11",
-    python_version = "3.11",
-)
-
 register_toolchains(
-    "@python_3_11_toolchains//:all",
+    "//envlogger:python3_system_toolchain",
+    "//envlogger:python3_system_cc_toolchain",
 )
 
+
+
 # Import Bazel protobuf rules.
 http_archive(
     name = "rules_proto",

diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -17,7 +17,7 @@ RUN apt update && apt install -y software-properties-common
 RUN add-apt-repository ppa:deadsnakes/ppa
 
 # Install necessary packages.
-RUN apt-get update && apt-get install -y git curl wget software-properties-common python3.$PY3_VERSION python3.$PY3_VERSION-dev libgmp-dev gcc-9 g++-9 tmux vim
+RUN apt-get update && apt-get install -y git curl wget software-properties-common python3.$PY3_VERSION python3.$PY3_VERSION-dev libgmp-dev gcc g++ tmux vim
 
 # Install distutils for Python 3.11 (distutils is removed in 3.12+).
 RUN if [ "$PY3_VERSION" = "11" ]; then apt install -y python3.11-distutils; fi
@@ -26,15 +26,14 @@ RUN if [ "$PY3_VERSION" = "11" ]; then apt install -y python3.11-distutils; fi
 RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.$PY3_VERSION
 
 # Download bazel.
-RUN wget https://github.com/bazelbuild/bazel/releases/download/6.2.1/bazel-6.2.1-linux-x86_64
-RUN chmod +x /bazel-6.2.1-linux-x86_64
-RUN mv /bazel-6.2.1-linux-x86_64 /usr/bin/bazel
+RUN wget https://github.com/bazelbuild/bazel/releases/download/7.1.0/bazel-7.1.0-linux-x86_64
+RUN chmod +x /bazel-7.1.0-linux-x86_64
+RUN mv /bazel-7.1.0-linux-x86_64 /usr/bin/bazel
 
 # Add python alternatives.
 RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.$PY3_VERSION 1
 
-# Override gcc/g++.
-RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-9 60 --slave /usr/bin/g++ g++ /usr/bin/g++-9
+
 
 # Install some basic things for all python versions.
 RUN echo 1 | update-alternatives --config python3

diff --git a/envlogger/BUILD.bazel b/envlogger/BUILD.bazel
@@ -14,6 +14,7 @@
 
 # Build targets for environment logger.
 load("@rules_python//python:defs.bzl", "py_library", "py_test")
+load("@rules_python//python/cc:py_cc_toolchain.bzl", "py_cc_toolchain")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -93,3 +94,23 @@ py_library(
     deps = [
     ],
 )
+
+# System Python toolchain for OSS build.
+toolchain(
+    name = "python3_system_toolchain",
+    toolchain = "@local_config_python//:py_runtime_pair",
+    toolchain_type = "@rules_python//python:toolchain_type",
+)
+
+py_cc_toolchain(
+    name = "python3_system_cc_toolchain_impl",
+    headers = "@local_config_python//:python_headers",
+    libs = "@local_config_python//:python_lib",
+    python_version = "3.11",
+)
+
+toolchain(
+    name = "python3_system_cc_toolchain",
+    toolchain = ":python3_system_cc_toolchain_impl",
+    toolchain_type = "@rules_python//python/cc:toolchain_type",
+)
diff --git a/envlogger/platform/default/proto_testutil.h b/envlogger/platform/default/proto_testutil.h
@@ -15,9 +15,12 @@
 #ifndef THIRD_PARTY_PY_ENVLOGGER_PLATFORM_DEFAULT_PROTO_TESTUTIL_H_
 #define THIRD_PARTY_PY_ENVLOGGER_PLATFORM_DEFAULT_PROTO_TESTUTIL_H_
 
+#include <memory>
+#include <ostream>
 #include <vector>
 
 #include "glog/logging.h"
+#include "google/protobuf/message.h"  // IWYU pragma: keep
 #include "google/protobuf/text_format.h"
 #include "google/protobuf/util/message_differencer.h"
 #include "gmock/gmock.h"
@@ -33,16 +36,30 @@ class ProtoStringMatcher {
  public:
   explicit ProtoStringMatcher(const std::string& expected)
       : expected_proto_str_(expected) {}
-  explicit ProtoStringMatcher(const google::protobuf::Message& expected)
-      : expected_proto_str_(expected.DebugString()) {}
+  explicit ProtoStringMatcher(const google::protobuf::Message& expected) {  // NOLINT
+    google::protobuf::Message* copy = expected.New();
+    copy->CopyFrom(expected);
+    expected_proto_ = std::shared_ptr<const google::protobuf::Message>(copy);
+  }
 
   template <typename Message>
   bool MatchAndExplain(const Message& actual_proto,
                        ::testing::MatchResultListener* listener) const;
 
-  void DescribeTo(::std::ostream* os) const { *os << expected_proto_str_; }
+  void DescribeTo(::std::ostream* os) const {
+    if (expected_proto_) {
+      *os << expected_proto_->ShortDebugString();
+    } else {
+      *os << expected_proto_str_;
+    }
+  }
   void DescribeNegationTo(::std::ostream* os) const {
-    *os << "not equal to expected message: " << expected_proto_str_;
+    *os << "not equal to expected message: ";
+    if (expected_proto_) {
+      *os << expected_proto_->ShortDebugString();
+    } else {
+      *os << expected_proto_str_;
+    }
   }
 
   void SetComparePartially() {
@@ -51,6 +68,7 @@ class ProtoStringMatcher {
 
  private:
   const std::string expected_proto_str_;
+  std::shared_ptr<const google::protobuf::Message> expected_proto_;
   google::protobuf::util::MessageDifferencer::Scope scope_ =
       google::protobuf::util::MessageDifferencer::FULL;
 };
@@ -66,7 +84,12 @@ template <typename Message>
 bool ProtoStringMatcher::MatchAndExplain(
     const Message& actual_proto,
     ::testing::MatchResultListener* listener) const {
-  Message expected_proto = CreateProto<Message>(expected_proto_str_);
+  Message expected_proto;
+  if (expected_proto_) {
+    expected_proto.CopyFrom(*expected_proto_);
+  } else {
+    expected_proto = CreateProto<Message>(expected_proto_str_);
+  }
 
   google::protobuf::util::MessageDifferencer differencer;
   std::string differences;

diff --git a/envlogger/requirements.txt b/envlogger/requirements.txt
@@ -1,6 +1,6 @@
 absl-py
 dm_env
 numpy
-protobuf
+protobuf>=3.14,<7.0.0
 tensorflow
 tfds-nightly