JumanjiWrapper

torchrl.envs.libs.jumanji.JumanjiWrapper(*args, **kwargs)

Jumanji environment wrapper.

Examples

>>> env = jumanji.make("Snake-6x6-v0")
>>> env = JumanjiWrapper(env)
>>> env.set_seed(0)
>>> td = env.reset()
>>> td["action"] = env.action_spec.rand()
>>> td = env.step(td)
>>> print(td1)
TensorDict(
    fields={
        action: Tensor(torch.Size([1]), dtype=torch.int32),
        done: Tensor(torch.Size([1]), dtype=torch.bool),
        next: TensorDict(
            fields={
                observation: Tensor(torch.Size([6, 6, 5]), dtype=torch.float32)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False),
        observation: Tensor(torch.Size([6, 6, 5]), dtype=torch.float32),
        reward: Tensor(torch.Size([1]), dtype=torch.float32),
        state: TensorDict(...)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)
>>> print(env.available_envs)
['Snake-6x6-v0', 'Snake-12x12-v0', 'TSP50-v0', 'TSP100-v0', ...]