Update

[ghstack-poisoned]

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

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)

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:
 #