From 97ab2e1e99cee04f5f2aae3f66e1a1311853812a Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 13 Nov 2024 16:58:18 +0000
Subject: [PATCH] Update

[ghstack-poisoned]
---
 docs/source/conf.py                  |  3 +++
 torchrl/_utils.py                    | 12 +++++++++++-
 tutorials/sphinx-tutorials/export.py |  6 ++++++
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/docs/source/conf.py b/docs/source/conf.py
index 0a7781b8675..35f5e5c3882 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -49,6 +49,8 @@
     version = f"main ({torchrl.__version__})"
     release = "main"
 
+os.environ["TORCHRL_CONSOLE_STREAM"] = "stdout"
+
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.
 #
@@ -95,6 +97,7 @@
     "abort_on_example_error": False,
     "only_warn_on_example_error": True,
     "show_memory": True,
+    "capture_repr": ("_repr_html_", "__repr__"),  # capture representations
 }
 
 napoleon_use_ivar = True
diff --git a/torchrl/_utils.py b/torchrl/_utils.py
index 31e00614fd9..e2f375e09ce 100644
--- a/torchrl/_utils.py
+++ b/torchrl/_utils.py
@@ -40,7 +40,17 @@
 # Remove all attached handlers
 while logger.hasHandlers():
     logger.removeHandler(logger.handlers[0])
-console_handler = logging.StreamHandler()
+stream_handlers = {
+    "stdout": sys.stdout,
+    "stderr": sys.stderr,
+}
+TORCHRL_CONSOLE_STREAM = os.getenv("TORCHRL_CONSOLE_STREAM")
+if TORCHRL_CONSOLE_STREAM:
+    stream_handler = stream_handlers[TORCHRL_CONSOLE_STREAM]
+else:
+    stream_handler = None
+console_handler = logging.StreamHandler(stream=stream_handler)
+
 console_handler.setLevel(logging.INFO)
 formatter = logging.Formatter("%(asctime)s [%(name)s][%(levelname)s] %(message)s")
 console_handler.setFormatter(formatter)
diff --git a/tutorials/sphinx-tutorials/export.py b/tutorials/sphinx-tutorials/export.py
index 0694621c99d..009ed686ab8 100644
--- a/tutorials/sphinx-tutorials/export.py
+++ b/tutorials/sphinx-tutorials/export.py
@@ -407,6 +407,12 @@
 screen_obs = ale.getScreenRGB()
 print("Observation from ALE simulator:", type(screen_obs), screen_obs.shape)
 
+from matplotlib import pyplot as plt
+
+plt.tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)
+plt.imshow(screen_obs)
+plt.title("Screen rendering of Pong game.")
+
 #####################################
 # Exporting to ONNX is quite similar the Export/AOTI above:
 #