Skip to content

Commit

Permalink
add comment on expected render() output
Browse files Browse the repository at this point in the history
  • Loading branch information
NickLucche committed Apr 25, 2024
1 parent 3d6acff commit 6b9d887
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions docs/guide/tensorboard.rst
Original file line number Diff line number Diff line change
Expand Up @@ -227,6 +227,9 @@ Here is an example of how to render an episode and log the resulting video to Te
:param _locals: A dictionary containing all local variables of the callback's scope
:param _globals: A dictionary containing all global variables of the callback's scope
"""
# We expect `render()` to return a uint8 array with values in [0, 255] or a float array
# with values in [0, 1], as described in
# https://pytorch.org/docs/stable/tensorboard.html#torch.utils.tensorboard.writer.SummaryWriter.add_video
screen = self._eval_env.render(mode="rgb_array")
# PyTorch uses CxHxW vs HxWxC gym (and tensorflow) image convention
screens.append(screen.transpose(2, 0, 1))
Expand Down

0 comments on commit 6b9d887

Please sign in to comment.