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)