JumanjiWrapper¶
- torchrl.envs.JumanjiWrapper(*args, **kwargs)[源代码]¶
Jumanji 的环境包装器。
Jumanji 提供了一个基于 Jax 的向量化模拟框架。TorchRL 的包装器会为了 jax 到 torch 的转换而产生一些开销，但计算图仍然可以在模拟轨迹之上构建，从而允许通过 rollout 进行反向传播。
GitHub: https://github.com/instadeepai/jumanji
文档: https://instadeepai.github.io/jumanji/
论文: https://arxiv.org/abs/2306.09884
注意
为了获得更好的性能，在实例化此类时请打开 jit。 jit 属性也可以在代码执行期间进行切换。
>>> env.jit = True # Used jit >>> env.jit = False # eager
- 参数:
env (jumanji.env.Environment) – 要包装的环境。
categorical_action_encoding (bool, optional) – 如果为
True，则分类规范将转换为 TorchRL 等效项（torchrl.data.Categorical），否则将使用独热编码（torchrl.data.OneHot）。默认为False。
- 关键字参数:
batch_size (torch.Size, optional) –
环境的批次大小。使用
jumanji时，这表示向量化环境的数量。如果批次大小为空，则环境未锁定批次，可以同时执行任意数量的环境。默认为torch.Size([])。>>> import jumanji >>> from torchrl.envs import JumanjiWrapper >>> base_env = jumanji.make("Snake-v1") >>> env = JumanjiWrapper(base_env) >>> # Set the batch-size of the TensorDict instead of the env allows to control the number >>> # of envs being run simultaneously >>> tdreset = env.reset(TensorDict(batch_size=[32])) >>> # Execute a rollout until all envs are done or max steps is reached, whichever comes first >>> rollout = env.rollout(100, break_when_all_done=True, auto_reset=False, tensordict=tdreset)
from_pixels (bool, optional) – 环境是否应渲染其输出。这将极大地影响环境吞吐量。只有第一个环境将被渲染。有关更多信息，请参见
render()。默认为 False。frame_skip (int, optional) – 如果提供，表示同一个动作需要重复多少步。返回的观察是序列的最后一次观察，而奖励将是跨步骤的奖励总和。
device (torch.device, optional) – 如果提供，表示数据要投射到的设备。默认为
torch.device("cpu")。allow_done_after_reset (bool, optional) – 如果为
True，则允许环境在调用reset()后立即完成。默认为False。jit (bool, optional) – 步进和重置方法是否应包装在 jit 中。默认为
False。
- 变量:
available_envs – 可用于构建的环境
示例
>>> import jumanji >>> from torchrl.envs import JumanjiWrapper >>> base_env = jumanji.make("Snake-v1") >>> env = JumanjiWrapper(base_env) >>> env.set_seed(0) >>> td = env.reset() >>> td["action"] = env.action_spec.rand() >>> td = env.step(td) >>> print(td) TensorDict( fields={ action: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False), action_mask: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.bool, is_shared=False), done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False), grid: Tensor(shape=torch.Size([12, 12, 5]), device=cpu, dtype=torch.float32, is_shared=False), next: TensorDict( fields={ action_mask: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.bool, is_shared=False), done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False), grid: Tensor(shape=torch.Size([12, 12, 5]), device=cpu, dtype=torch.float32, is_shared=False), reward: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False), state: TensorDict( fields={ action_mask: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.bool, is_shared=False), body: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.bool, is_shared=False), body_state: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.int32, is_shared=False), fruit_position: TensorDict( fields={ col: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False), row: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False), head_position: TensorDict( fields={ col: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False), row: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False), key: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int32, is_shared=False), length: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False), step_count: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False), tail: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.bool, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False), step_count: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False), terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False), state: TensorDict( fields={ action_mask: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.bool, is_shared=False), body: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.bool, is_shared=False), body_state: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.int32, is_shared=False), fruit_position: TensorDict( fields={ col: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False), row: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False), head_position: TensorDict( fields={ col: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False), row: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False), key: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int32, is_shared=False), length: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False), step_count: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False), tail: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.bool, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False), step_count: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False), terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False) >>> print(env.available_envs) ['Game2048-v1', 'Maze-v0', 'Cleaner-v0', 'CVRP-v1', 'MultiCVRP-v0', 'Minesweeper-v0', 'RubiksCube-v0', 'Knapsack-v1', 'Sudoku-v0', 'Snake-v1', 'TSP-v1', 'Connector-v2', 'MMST-v0', 'GraphColoring-v0', 'RubiksCube-partly-scrambled-v0', 'RobotWarehouse-v0', 'Tetris-v0', 'BinPack-v2', 'Sudoku-very-easy-v0', 'JobShop-v0']
为了利用 Jumanji 的优势，通常会同时执行多个环境。
>>> import jumanji >>> from torchrl.envs import JumanjiWrapper >>> base_env = jumanji.make("Snake-v1") >>> env = JumanjiWrapper(base_env, batch_size=[10]) >>> env.set_seed(0) >>> td = env.reset() >>> td["action"] = env.action_spec.rand() >>> td = env.step(td)
在下面的示例中，我们迭代测试不同的批次大小，并报告短 rollout 的执行时间。
示例
>>> from torch.utils.benchmark import Timer >>> for batch_size in [4, 16, 128]: ... timer = Timer( ... ''' ... env.rollout(100) ... ''', ... setup=f''' ... from torchrl.envs import JumanjiWrapper ... import jumanji ... env = JumanjiWrapper(jumanji.make('Snake-v1'), batch_size=[{batch_size}]) ... env.set_seed(0) ... env.rollout(2) ... ''') ... print(batch_size, timer.timeit(number=10)) 4 env.rollout(100) setup: [...] Median: 122.40 ms 2 measurements, 1 runs per measurement, 1 thread
16 env.rollout(100) 设置： […] 中位数：134.39 毫秒 2 次测量，每次测量 1 次运行，1 个线程
128 env.rollout(100) 设置： […] 中位数：172.31 毫秒 2 次测量，每次测量 1 次运行，1 个线程
