torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

sota-implementations/a2c/utils_atari.py

@@ -0,0 +1,240 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import numpy as np
+import torch.nn
+import torch.optim
+from tensordict.nn import TensorDictModule
+from torchrl.data.tensor_specs import CategoricalBox
+from torchrl.envs import (
+    CatFrames,
+    DoubleToFloat,
+    EndOfLifeTransform,
+    EnvCreator,
+    ExplorationType,
+    GrayScale,
+    GymEnv,
+    NoopResetEnv,
+    ParallelEnv,
+    Resize,
+    RewardSum,
+    set_gym_backend,
+    SignTransform,
+    StepCounter,
+    ToTensorImage,
+    TransformedEnv,
+    VecNorm,
+)
+from torchrl.modules import (
+    ActorValueOperator,
+    ConvNet,
+    MLP,
+    OneHotCategorical,
+    ProbabilisticActor,
+    TanhNormal,
+    ValueOperator,
+)
+from torchrl.record import VideoRecorder
+
+
+# ====================================================================
+# Environment utils
+# --------------------------------------------------------------------
+
+
+def make_base_env(
+    env_name="BreakoutNoFrameskip-v4",
+    gym_backend="gymnasium",
+    frame_skip=4,
+    device="cpu",
+    is_test=False,
+):
+    with set_gym_backend(gym_backend):
+        env = GymEnv(
+            env_name,
+            frame_skip=frame_skip,
+            from_pixels=True,
+            pixels_only=False,
+            device=device,
+        )
+    env = TransformedEnv(env)
+    env.append_transform(NoopResetEnv(noops=30, random=True))
+    if not is_test:
+        env.append_transform(EndOfLifeTransform())
+    return env
+
+
+def make_parallel_env(env_name, num_envs, device, gym_backend, is_test=False):
+    env = ParallelEnv(
+        num_envs,
+        EnvCreator(
+            lambda: make_base_env(env_name, gym_backend=gym_backend, is_test=is_test),
+        ),
+        serial_for_single=True,
+        device=device,
+    )
+    env = TransformedEnv(env)
+    env.append_transform(DoubleToFloat())
+    env.append_transform(ToTensorImage())
+    env.append_transform(GrayScale())
+    env.append_transform(Resize(84, 84))
+    env.append_transform(CatFrames(N=4, dim=-3))
+    env.append_transform(RewardSum())
+    env.append_transform(StepCounter(max_steps=4500))
+    if not is_test:
+        env.append_transform(SignTransform(in_keys=["reward"]))
+    env.append_transform(VecNorm(in_keys=["pixels"]))
+    return env
+
+
+# ====================================================================
+# Model utils
+# --------------------------------------------------------------------
+
+
+def make_ppo_modules_pixels(proof_environment, device):
+
+    # Define input shape
+    input_shape = proof_environment.observation_spec["pixels"].shape
+
+    # Define distribution class and kwargs
+    if isinstance(proof_environment.action_spec_unbatched.space, CategoricalBox):
+        num_outputs = proof_environment.action_spec_unbatched.space.n
+        distribution_class = OneHotCategorical
+        distribution_kwargs = {}
+    else:  # is ContinuousBox
+        num_outputs = proof_environment.action_spec_unbatched.shape
+        distribution_class = TanhNormal
+        distribution_kwargs = {
+            "low": proof_environment.action_spec_unbatched.space.low.to(device),
+            "high": proof_environment.action_spec_unbatched.space.high.to(device),
+        }
+
+    # Define input keys
+    in_keys = ["pixels"]
+
+    # Define a shared Module and TensorDictModule (CNN + MLP)
+    common_cnn = ConvNet(
+        activation_class=torch.nn.ReLU,
+        num_cells=[32, 64, 64],
+        kernel_sizes=[8, 4, 3],
+        strides=[4, 2, 1],
+        device=device,
+    )
+    common_cnn_output = common_cnn(torch.ones(input_shape, device=device))
+    common_mlp = MLP(
+        in_features=common_cnn_output.shape[-1],
+        activation_class=torch.nn.ReLU,
+        activate_last_layer=True,
+        out_features=512,
+        num_cells=[],
+        device=device,
+    )
+    common_mlp_output = common_mlp(common_cnn_output)
+
+    # Define shared net as TensorDictModule
+    common_module = TensorDictModule(
+        module=torch.nn.Sequential(common_cnn, common_mlp),
+        in_keys=in_keys,
+        out_keys=["common_features"],
+    )
+
+    # Define on head for the policy
+    policy_net = MLP(
+        in_features=common_mlp_output.shape[-1],
+        out_features=num_outputs,
+        activation_class=torch.nn.ReLU,
+        num_cells=[],
+        device=device,
+    )
+    policy_module = TensorDictModule(
+        module=policy_net,
+        in_keys=["common_features"],
+        out_keys=["logits"],
+    )
+
+    # Add probabilistic sampling of the actions
+    policy_module = ProbabilisticActor(
+        policy_module,
+        in_keys=["logits"],
+        spec=proof_environment.full_action_spec_unbatched.to(device),
+        distribution_class=distribution_class,
+        distribution_kwargs=distribution_kwargs,
+        return_log_prob=True,
+        default_interaction_type=ExplorationType.RANDOM,
+    )
+
+    # Define another head for the value
+    value_net = MLP(
+        activation_class=torch.nn.ReLU,
+        in_features=common_mlp_output.shape[-1],
+        out_features=1,
+        num_cells=[],
+        device=device,
+    )
+    value_module = ValueOperator(
+        value_net,
+        in_keys=["common_features"],
+    )
+
+    return common_module, policy_module, value_module
+
+
+def make_ppo_models(env_name, device, gym_backend):
+
+    proof_environment = make_parallel_env(
+        env_name, num_envs=1, device="cpu", gym_backend=gym_backend
+    )
+    common_module, policy_module, value_module = make_ppo_modules_pixels(
+        proof_environment, device=device
+    )
+
+    # Wrap modules in a single ActorCritic operator
+    actor_critic = ActorValueOperator(
+        common_operator=common_module,
+        policy_operator=policy_module,
+        value_operator=value_module,
+    )
+
+    with torch.no_grad():
+        td = proof_environment.fake_tensordict().expand(1)
+        td = actor_critic(td.to(device))
+        del td
+
+    actor = actor_critic.get_policy_operator()
+    critic = actor_critic.get_value_operator()
+    critic_head = actor_critic.get_value_head()
+
+    del proof_environment
+
+    return actor, critic, critic_head
+
+
+# ====================================================================
+# Evaluation utils
+# --------------------------------------------------------------------
+
+
+def dump_video(module):
+    if isinstance(module, VideoRecorder):
+        module.dump()
+
+
+def eval_model(actor, test_env, num_episodes=3):
+    test_rewards = []
+    for _ in range(num_episodes):
+        td_test = test_env.rollout(
+            policy=actor,
+            auto_reset=True,
+            auto_cast_to_device=True,
+            break_when_any_done=True,
+            max_steps=10_000_000,
+        )
+        reward = td_test["next", "episode_reward"][td_test["next", "done"]]
+        test_rewards = np.append(test_rewards, reward.cpu().numpy())
+        test_env.apply(dump_video)
+    del td_test
+    return test_rewards.mean()

sota-implementations/a2c/utils_mujoco.py

@@ -0,0 +1,160 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import numpy as np
+import torch.nn
+import torch.optim
+
+from tensordict.nn import AddStateIndependentNormalScale, TensorDictModule
+from torchrl.envs import (
+    ClipTransform,
+    DoubleToFloat,
+    ExplorationType,
+    RewardSum,
+    StepCounter,
+    TransformedEnv,
+    VecNorm,
+)
+from torchrl.envs.libs.gym import GymEnv
+from torchrl.modules import MLP, ProbabilisticActor, TanhNormal, ValueOperator
+from torchrl.record import VideoRecorder
+
+
+# ====================================================================
+# Environment utils
+# --------------------------------------------------------------------
+
+
+def make_env(
+    env_name="HalfCheetah-v4", device="cpu", from_pixels=False, pixels_only=False
+):
+    env = GymEnv(
+        env_name, device=device, from_pixels=from_pixels, pixels_only=pixels_only
+    )
+    env = TransformedEnv(env)
+    env.append_transform(RewardSum())
+    env.append_transform(StepCounter())
+    env.append_transform(VecNorm(in_keys=["observation"]))
+    env.append_transform(ClipTransform(in_keys=["observation"], low=-10, high=10))
+    env.append_transform(DoubleToFloat(in_keys=["observation"]))
+    return env
+
+
+# ====================================================================
+# Model utils
+# --------------------------------------------------------------------
+
+
+def make_ppo_models_state(proof_environment, device, *, compile: bool = False):
+
+    # Define input shape
+    input_shape = proof_environment.observation_spec["observation"].shape
+
+    # Define policy output distribution class
+    num_outputs = proof_environment.action_spec_unbatched.shape[-1]
+    distribution_class = TanhNormal
+    distribution_kwargs = {
+        "low": proof_environment.action_spec_unbatched.space.low.to(device),
+        "high": proof_environment.action_spec_unbatched.space.high.to(device),
+        "tanh_loc": False,
+        "safe_tanh": True,
+    }
+
+    # Define policy architecture
+    policy_mlp = MLP(
+        in_features=input_shape[-1],
+        activation_class=torch.nn.Tanh,
+        out_features=num_outputs,  # predict only loc
+        num_cells=[64, 64],
+        device=device,
+    )
+
+    # Initialize policy weights
+    for layer in policy_mlp.modules():
+        if isinstance(layer, torch.nn.Linear):
+            torch.nn.init.orthogonal_(layer.weight, 1.0)
+            layer.bias.data.zero_()
+
+    # Add state-independent normal scale
+    policy_mlp = torch.nn.Sequential(
+        policy_mlp,
+        AddStateIndependentNormalScale(
+            proof_environment.action_spec_unbatched.shape[-1], device=device
+        ),
+    )
+
+    # Add probabilistic sampling of the actions
+    policy_module = ProbabilisticActor(
+        TensorDictModule(
+            module=policy_mlp,
+            in_keys=["observation"],
+            out_keys=["loc", "scale"],
+        ),
+        in_keys=["loc", "scale"],
+        spec=proof_environment.full_action_spec_unbatched.to(device),
+        distribution_class=distribution_class,
+        distribution_kwargs=distribution_kwargs,
+        return_log_prob=True,
+        default_interaction_type=ExplorationType.RANDOM,
+    )
+
+    # Define value architecture
+    value_mlp = MLP(
+        in_features=input_shape[-1],
+        activation_class=torch.nn.Tanh,
+        out_features=1,
+        num_cells=[64, 64],
+        device=device,
+    )
+
+    # Initialize value weights
+    for layer in value_mlp.modules():
+        if isinstance(layer, torch.nn.Linear):
+            torch.nn.init.orthogonal_(layer.weight, 0.01)
+            layer.bias.data.zero_()
+
+    # Define value module
+    value_module = ValueOperator(
+        value_mlp,
+        in_keys=["observation"],
+    )
+
+    return policy_module, value_module
+
+
+def make_ppo_models(env_name, device, *, compile: bool = False):
+    proof_environment = make_env(env_name, device="cpu")
+    actor, critic = make_ppo_models_state(
+        proof_environment, device=device, compile=compile
+    )
+    return actor, critic
+
+
+# ====================================================================
+# Evaluation utils
+# --------------------------------------------------------------------
+
+
+def dump_video(module):
+    if isinstance(module, VideoRecorder):
+        module.dump()
+
+
+def eval_model(actor, test_env, num_episodes=3):
+    test_rewards = []
+    for _ in range(num_episodes):
+        td_test = test_env.rollout(
+            policy=actor,
+            auto_reset=True,
+            auto_cast_to_device=True,
+            break_when_any_done=True,
+            max_steps=10_000_000,
+        )
+        reward = td_test["next", "episode_reward"][td_test["next", "done"]]
+        test_rewards = np.append(test_rewards, reward.cpu().numpy())
+        test_env.apply(dump_video)
+    del td_test
+    return test_rewards.mean()

sota-implementations/bandits/README.md

@@ -0,0 +1,7 @@
+# Bandits example
+
+## Note:
+This example is not included in the benchmarked results of the current release (v0.3). The intention is to include it in the
+benchmarking of future releases, to ensure that it can be successfully run with the release code and that the
+results are consistent. For now, be aware that this additional check has not been performed in the case of this
+specific example.

sota-implementations/bandits/dqn.py

@@ -0,0 +1,126 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import argparse
+
+import torch
+import tqdm
+
+from tensordict.nn import TensorDictSequential
+from torch import nn
+from torchrl.envs.libs.openml import OpenMLEnv
+from torchrl.envs.utils import ExplorationType, set_exploration_type
+from torchrl.modules import DistributionalQValueActor, EGreedyModule, MLP, QValueActor
+from torchrl.objectives import DistributionalDQNLoss, DQNLoss
+
+parser = argparse.ArgumentParser()
+
+# Add arguments
+parser.add_argument("--batch_size", type=int, default=256, help="batch size")
+parser.add_argument("--n_steps", type=int, default=10000, help="number of steps")
+parser.add_argument(
+    "--eps_greedy", type=float, default=0.1, help="epsilon-greedy parameter"
+)
+parser.add_argument("--lr", type=float, default=2e-4, help="learning rate")
+parser.add_argument("--wd", type=float, default=1e-4, help="weight decay")
+parser.add_argument("--n_cells", type=int, default=128, help="number of cells")
+parser.add_argument(
+    "--distributional", action="store_true", help="enable distributional Q-learning"
+)
+parser.add_argument(
+    "--dataset",
+    default="adult_onehot",
+    choices=[
+        "adult_num",
+        "adult_onehot",
+        "mushroom_num",
+        "mushroom_onehot",
+        "covertype",
+        "shuttle",
+        "magic",
+    ],
+    help="OpenML dataset",
+)
+
+if __name__ == "__main__":
+    # Parse arguments
+    args = parser.parse_args()
+
+    # Access arguments
+    batch_size = args.batch_size
+    n_steps = args.n_steps
+    eps_greedy = args.eps_greedy
+    lr = args.lr
+    wd = args.wd
+    n_cells = args.n_cells
+    distributional = args.distributional
+    dataset = args.dataset
+
+    env = OpenMLEnv(dataset, batch_size=[batch_size])
+    n_actions = env.action_spec.space.n
+    if distributional:
+        # does not really make sense since the value is either 0 or 1 and hopefully we
+        # should always predict 1
+        nbins = 2
+        model = MLP(
+            out_features=(nbins, n_actions),
+            depth=3,
+            num_cells=n_cells,
+            activation_class=nn.Tanh,
+        )
+        actor = DistributionalQValueActor(
+            model, support=torch.arange(2), action_space="categorical"
+        )
+        actor(env.reset())
+        loss = DistributionalDQNLoss(
+            actor,
+        )
+        loss.make_value_estimator(gamma=0.9)
+    else:
+        model = MLP(
+            out_features=n_actions, depth=3, num_cells=n_cells, activation_class=nn.Tanh
+        )
+        actor = QValueActor(model, action_space="categorical")
+        actor(env.reset())
+        loss = DQNLoss(actor, loss_function="smooth_l1", action_space=env.action_spec)
+        loss.make_value_estimator(gamma=0.0)
+    policy = TensorDictSequential(
+        actor,
+        EGreedyModule(
+            eps_init=eps_greedy,
+            eps_end=0.0,
+            annealing_num_steps=n_steps,
+            spec=env.action_spec,
+        ),
+    )
+    optim = torch.optim.Adam(loss.parameters(), lr, weight_decay=wd)
+
+    pbar = tqdm.tqdm(range(n_steps))
+
+    init_r = None
+    init_loss = None
+    for i in pbar:
+        with set_exploration_type(ExplorationType.RANDOM):
+            data = env.step(policy(env.reset()))
+        loss_vals = loss(data)
+        loss_val = sum(
+            value for key, value in loss_vals.items() if key.startswith("loss")
+        )
+        loss_val.backward()
+        optim.step()
+        optim.zero_grad()
+        if i % 10 == 0:
+            test_data = env.step(policy(env.reset()))
+            if init_r is None:
+                init_r = test_data["next", "reward"].sum() / env.numel()
+            if init_loss is None:
+                init_loss = loss_val.detach().item()
+            pbar.set_description(
+                f"reward: {test_data['next', 'reward'].sum() / env.numel(): 4.4f} (init={init_r: 4.4f}), "
+                f"training reward {data['next', 'reward'].sum() / env.numel() : 4.4f}, "
+                f"loss {loss_val: 4.4f} (init: {init_loss: 4.4f})"
+            )
+        policy[1].step()