快捷方式

PixelRenderTransform

torchrl.record.PixelRenderTransform(out_keys: List[NestedKey] = None, preproc: Callable[[np.ndarray | torch.Tensor], np.ndarray | torch.Tensor] = None, as_non_tensor: bool = None, render_method: str = 'render', pass_tensordict: bool = False, **kwargs) None[source]

一个 transform,用于在其父环境上调用 render 方法,并在 tensordict 中注册像素观察结果。

这个 transform 提供了一种替代 from_pixels 语法糖的方式,适用于实例化支持渲染的环境成本较高,或 from_pixels 未实现的情况。它既可以在单个环境中使用,也可以在批量环境中类似地使用。

参数:
  • out_keys (List[NestedKey] 或 Nested) – 用于注册像素观察结果的键列表。

  • preproc (Callable, 可选) – 预处理函数。可用于重塑观察结果,或应用任何其他使其能够在输出数据中注册的转换。

  • as_non_tensor (bool, 可选) – 如果为 True,数据将作为 NonTensorData 写入,从而放宽形状要求。如果未提供,将根据输入数据类型和形状自动推断。

  • render_method (str, 可选) – render 方法的名称。默认为 "render"

  • pass_tensordict (bool, 可选) – 如果为 True,输入的 tensordict 将传递给 render 方法。这使得无状态环境能够进行渲染。默认为 False

  • **kwargs – 要传递给 render 函数的额外关键字参数(例如 mode="rgb_array")。

示例

>>> from torchrl.envs import GymEnv, check_env_specs, ParallelEnv, EnvCreator
>>> from torchrl.record.loggers import CSVLogger
>>> from torchrl.record.recorder import PixelRenderTransform, VideoRecorder
>>>
>>> def make_env():
>>>     env = GymEnv("CartPole-v1", render_mode="rgb_array")
>>>     env = env.append_transform(PixelRenderTransform())
>>>     return env
>>>
>>> if __name__ == "__main__":
...     logger = CSVLogger("dummy", video_format="mp4")
...
...     env = ParallelEnv(4, EnvCreator(make_env))
...
...     env = env.append_transform(VideoRecorder(logger=logger, tag="pixels_record"))
...     env.rollout(3)
...
...     check_env_specs(env)
...
...     r = env.rollout(30)
...     print(env)
...     env.transform.dump()
...     env.close()

当批量环境的 render() 方法返回单个图像时,也可以使用此 transform

示例

>>> from torchrl.envs import check_env_specs
>>> from torchrl.envs.libs.vmas import VmasEnv
>>> from torchrl.record.loggers import CSVLogger
>>> from torchrl.record.recorder import PixelRenderTransform, VideoRecorder
>>>
>>> env = VmasEnv(
...     scenario="flocking",
...     num_envs=32,
...     continuous_actions=True,
...     max_steps=200,
...     device="cpu",
...     seed=None,
...     # Scenario kwargs
...     n_agents=5,
... )
>>>
>>> logger = CSVLogger("dummy", video_format="mp4")
>>>
>>> env = env.append_transform(PixelRenderTransform(mode="rgb_array", preproc=lambda x: x.copy()))
>>> env = env.append_transform(VideoRecorder(logger=logger, tag="pixels_record"))
>>>
>>> check_env_specs(env)
>>>
>>> r = env.rollout(30)
>>> env.transform[-1].dump()

可以使用 switch() 方法禁用 transform,这将使渲染在关闭时打开,在打开时关闭(也可以传递参数来控制此行为)。由于 transforms 是 Module 实例,可以使用 apply() 来控制此行为

>>> def switch(module):
...     if isinstance(module, PixelRenderTransform):
...         module.switch()
>>> env.apply(switch)

文档

获取 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深度教程

查看教程

资源

查找开发资源并获取问题解答

查看资源