torchrl 0.11.0__cp314-cp314t-win_amd64.whl

sota-implementations/impala/impala_single_node.py

@@ -0,0 +1,261 @@
+# 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.
+
+"""
+This script reproduces the IMPALA Algorithm
+results from Espeholt et al. 2018 for the on Atari Environments.
+"""
+from __future__ import annotations
+
+import hydra
+from torchrl._utils import logger as torchrl_logger
+
+
+@hydra.main(config_path="", config_name="config_single_node", version_base="1.1")
+def main(cfg: DictConfig):  # noqa: F821
+
+    import time
+
+    import torch.optim
+    import tqdm
+
+    from tensordict import TensorDict
+    from torchrl.collectors import MultiaSyncDataCollector
+    from torchrl.data import LazyMemmapStorage, TensorDictReplayBuffer
+    from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
+    from torchrl.envs import ExplorationType, set_exploration_type
+    from torchrl.objectives import A2CLoss
+    from torchrl.objectives.value import VTrace
+    from torchrl.record.loggers import generate_exp_name, get_logger
+    from utils import eval_model, make_env, make_ppo_models
+
+    device = cfg.device
+    if not device:
+        device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0")
+    else:
+        device = torch.device(device)
+
+    # Correct for frame_skip
+    frame_skip = 4
+    total_frames = cfg.collector.total_frames // frame_skip
+    frames_per_batch = cfg.collector.frames_per_batch // frame_skip
+    test_interval = cfg.logger.test_interval // frame_skip
+
+    # Extract other config parameters
+    batch_size = cfg.loss.batch_size  # Number of rollouts per batch
+    num_workers = (
+        cfg.collector.num_workers
+    )  # Number of parallel workers collecting rollouts
+    lr = cfg.optim.lr
+    anneal_lr = cfg.optim.anneal_lr
+    sgd_updates = cfg.loss.sgd_updates
+    max_grad_norm = cfg.optim.max_grad_norm
+    num_test_episodes = cfg.logger.num_test_episodes
+    total_network_updates = (
+        total_frames // (frames_per_batch * batch_size)
+    ) * cfg.loss.sgd_updates
+
+    # Create models (check utils.py)
+    actor, critic = make_ppo_models(cfg.env.env_name, cfg.env.backend)
+
+    # Create collector
+    collector = MultiaSyncDataCollector(
+        create_env_fn=[make_env(cfg.env.env_name, device, gym_backend=cfg.env.backend)]
+        * num_workers,
+        policy=actor,
+        frames_per_batch=frames_per_batch,
+        total_frames=total_frames,
+        device=device,
+        storing_device=device,
+        max_frames_per_traj=-1,
+        update_at_each_batch=True,
+    )
+
+    # Create data buffer
+    sampler = SamplerWithoutReplacement()
+    data_buffer = TensorDictReplayBuffer(
+        storage=LazyMemmapStorage(frames_per_batch * batch_size),
+        sampler=sampler,
+        batch_size=frames_per_batch * batch_size,
+    )
+
+    # Create loss and adv modules
+    adv_module = VTrace(
+        gamma=cfg.loss.gamma,
+        value_network=critic,
+        actor_network=actor,
+        average_adv=False,
+    )
+    loss_module = A2CLoss(
+        actor_network=actor,
+        critic_network=critic,
+        loss_critic_type=cfg.loss.loss_critic_type,
+        entropy_coeff=cfg.loss.entropy_coeff,
+        critic_coeff=cfg.loss.critic_coeff,
+    )
+    loss_module.set_keys(done="eol", terminated="eol")
+
+    # Create optimizer
+    optim = torch.optim.RMSprop(
+        loss_module.parameters(),
+        lr=cfg.optim.lr,
+        weight_decay=cfg.optim.weight_decay,
+        eps=cfg.optim.eps,
+        alpha=cfg.optim.alpha,
+    )
+
+    # Create logger
+    logger = None
+    if cfg.logger.backend:
+        exp_name = generate_exp_name(
+            "IMPALA", f"{cfg.logger.exp_name}_{cfg.env.env_name}"
+        )
+        logger = get_logger(
+            cfg.logger.backend,
+            logger_name="impala",
+            experiment_name=exp_name,
+            wandb_kwargs={
+                "config": dict(cfg),
+                "project": cfg.logger.project_name,
+                "group": cfg.logger.group_name,
+            },
+        )
+
+    # Create test environment
+    test_env = make_env(
+        cfg.env.env_name, device, gym_backend=cfg.env.backend, is_test=True
+    )
+    test_env.eval()
+
+    # Main loop
+    collected_frames = 0
+    num_network_updates = 0
+    pbar = tqdm.tqdm(total=total_frames)
+    accumulator = []
+    start_time = sampling_start = time.time()
+    for i, data in enumerate(collector):
+
+        metrics_to_log = {}
+        sampling_time = time.time() - sampling_start
+        frames_in_batch = data.numel()
+        collected_frames += frames_in_batch * frame_skip
+        pbar.update(data.numel())
+
+        # Get training rewards and episode lengths
+        episode_rewards = data["next", "episode_reward"][data["next", "terminated"]]
+        if len(episode_rewards) > 0:
+            episode_length = data["next", "step_count"][data["next", "terminated"]]
+            metrics_to_log.update(
+                {
+                    "train/reward": episode_rewards.mean().item(),
+                    "train/episode_length": episode_length.sum().item()
+                    / len(episode_length),
+                }
+            )
+
+        if len(accumulator) < batch_size:
+            accumulator.append(data)
+            if logger:
+                for key, value in metrics_to_log.items():
+                    logger.log_scalar(key, value, collected_frames)
+            continue
+
+        losses = TensorDict(batch_size=[sgd_updates])
+        training_start = time.time()
+        for j in range(sgd_updates):
+
+            # Create a single batch of trajectories
+            stacked_data = torch.stack(accumulator, dim=0).contiguous()
+            stacked_data = stacked_data.to(device, non_blocking=True)
+
+            # Compute advantage
+            with torch.no_grad():
+                stacked_data = adv_module(stacked_data)
+
+            # Add to replay buffer
+            for stacked_d in stacked_data:
+                stacked_data_reshape = stacked_d.reshape(-1)
+                data_buffer.extend(stacked_data_reshape)
+
+            for batch in data_buffer:
+
+                # Linearly decrease the learning rate and clip epsilon
+                alpha = 1.0
+                if anneal_lr:
+                    alpha = 1 - (num_network_updates / total_network_updates)
+                    for group in optim.param_groups:
+                        group["lr"] = lr * alpha
+                num_network_updates += 1
+
+                # Get a data batch
+                batch = batch.to(device, non_blocking=True)
+
+                # Forward pass loss
+                loss = loss_module(batch)
+                losses[j] = loss.select(
+                    "loss_critic", "loss_entropy", "loss_objective"
+                ).detach()
+                loss_sum = (
+                    loss["loss_critic"] + loss["loss_objective"] + loss["loss_entropy"]
+                )
+
+                # Backward pass
+                loss_sum.backward()
+                torch.nn.utils.clip_grad_norm_(
+                    list(loss_module.parameters()), max_norm=max_grad_norm
+                )
+
+                # Update the networks
+                optim.step()
+                optim.zero_grad()
+
+        # Get training losses and times
+        training_time = time.time() - training_start
+        losses = losses.apply(lambda x: x.float().mean(), batch_size=[])
+        for key, value in losses.items():
+            metrics_to_log.update({f"train/{key}": value.item()})
+        metrics_to_log.update(
+            {
+                "train/lr": alpha * lr,
+                "train/sampling_time": sampling_time,
+                "train/training_time": training_time,
+            }
+        )
+
+        # Get test rewards
+        with torch.no_grad(), set_exploration_type(ExplorationType.DETERMINISTIC):
+            if ((i - 1) * frames_in_batch * frame_skip) // test_interval < (
+                i * frames_in_batch * frame_skip
+            ) // test_interval:
+                actor.eval()
+                eval_start = time.time()
+                test_reward = eval_model(
+                    actor, test_env, num_episodes=num_test_episodes
+                )
+                eval_time = time.time() - eval_start
+                metrics_to_log.update(
+                    {
+                        "eval/reward": test_reward,
+                        "eval/time": eval_time,
+                    }
+                )
+                actor.train()
+
+        if logger:
+            for key, value in metrics_to_log.items():
+                logger.log_scalar(key, value, collected_frames)
+
+        collector.update_policy_weights_()
+        sampling_start = time.time()
+        accumulator = []
+
+    collector.shutdown()
+    end_time = time.time()
+    execution_time = end_time - start_time
+    torchrl_logger.info(f"Training took {execution_time:.2f} seconds to finish")
+
+
+if __name__ == "__main__":
+    main()

sota-implementations/impala/utils.py

@@ -0,0 +1,184 @@
+# 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 torch.nn
+import torch.optim
+from tensordict.nn import TensorDictModule
+from torchrl.envs import (
+    CatFrames,
+    DoubleToFloat,
+    EndOfLifeTransform,
+    ExplorationType,
+    GrayScale,
+    GymEnv,
+    NoopResetEnv,
+    Resize,
+    RewardSum,
+    set_gym_backend,
+    SignTransform,
+    StepCounter,
+    ToTensorImage,
+    TransformedEnv,
+    VecNorm,
+)
+from torchrl.modules import (
+    ActorValueOperator,
+    ConvNet,
+    MLP,
+    OneHotCategorical,
+    ProbabilisticActor,
+    ValueOperator,
+)
+
+
+# ====================================================================
+# Environment utils
+# --------------------------------------------------------------------
+
+
+def make_env(env_name, device, gym_backend, is_test=False):
+    with set_gym_backend(gym_backend):
+        env = GymEnv(
+            env_name, frame_skip=4, 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())
+        env.append_transform(SignTransform(in_keys=["reward"]))
+    env.append_transform(ToTensorImage(from_int=False))
+    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))
+    env.append_transform(DoubleToFloat())
+    env.append_transform(VecNorm(in_keys=["pixels"]))
+    return env
+
+
+# ====================================================================
+# Model utils
+# --------------------------------------------------------------------
+
+
+def make_ppo_modules_pixels(proof_environment):
+
+    # Define input shape
+    input_shape = proof_environment.observation_spec["pixels"].shape
+
+    # Define distribution class and kwargs
+    num_outputs = proof_environment.action_spec_unbatched.space.n
+    distribution_class = OneHotCategorical
+    distribution_kwargs = {}
+
+    # 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],
+    )
+    common_cnn_output = common_cnn(torch.ones(input_shape))
+    common_mlp = MLP(
+        in_features=common_cnn_output.shape[-1],
+        activation_class=torch.nn.ReLU,
+        activate_last_layer=True,
+        out_features=512,
+        num_cells=[],
+    )
+    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 one 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=[],
+    )
+    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,
+        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=[],
+    )
+    value_module = ValueOperator(
+        value_net,
+        in_keys=["common_features"],
+    )
+
+    return common_module, policy_module, value_module
+
+
+def make_ppo_models(env_name, gym_backend):
+
+    proof_environment = make_env(env_name, device="cpu", gym_backend=gym_backend)
+    common_module, policy_module, value_module = make_ppo_modules_pixels(
+        proof_environment
+    )
+
+    # Wrap modules in a single ActorCritic operator
+    actor_critic = ActorValueOperator(
+        common_operator=common_module,
+        policy_operator=policy_module,
+        value_operator=value_module,
+    )
+
+    actor = actor_critic.get_policy_operator()
+    critic = actor_critic.get_value_operator()
+
+    del proof_environment
+
+    return actor, critic
+
+
+# ====================================================================
+# Evaluation utils
+# --------------------------------------------------------------------
+
+
+def eval_model(actor, test_env, num_episodes=3):
+    test_rewards = torch.zeros(num_episodes, dtype=torch.float32)
+    for i 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[i] = reward.sum()
+    del td_test
+    return test_rewards.mean()

sota-implementations/iql/discrete_iql.py

@@ -0,0 +1,230 @@
+# 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.
+"""IQL Example.
+
+This is a self-contained example of an online discrete IQL training script.
+
+It works across Gym and MuJoCo over a variety of tasks.
+
+The helper functions are coded in the utils.py associated with this script.
+
+"""
+from __future__ import annotations
+
+import warnings
+
+import hydra
+import numpy as np
+import torch
+import tqdm
+from tensordict import TensorDict
+from tensordict.nn import CudaGraphModule
+from torchrl._utils import get_available_device, timeit
+from torchrl.envs import set_gym_backend
+from torchrl.envs.utils import ExplorationType, set_exploration_type
+from torchrl.objectives import group_optimizers
+from torchrl.record.loggers import generate_exp_name, get_logger
+from utils import (
+    dump_video,
+    log_metrics,
+    make_collector,
+    make_discrete_iql_model,
+    make_discrete_loss,
+    make_environment,
+    make_iql_optimizer,
+    make_replay_buffer,
+)
+
+torch.set_float32_matmul_precision("high")
+
+
+@hydra.main(config_path="", config_name="discrete_iql")
+def main(cfg: DictConfig):  # noqa: F821
+    set_gym_backend(cfg.env.backend).set()
+
+    # Create logger
+    exp_name = generate_exp_name("Discrete-IQL-online", cfg.logger.exp_name)
+    logger = None
+    if cfg.logger.backend:
+        logger = get_logger(
+            logger_type=cfg.logger.backend,
+            logger_name="iql_logging",
+            experiment_name=exp_name,
+            wandb_kwargs={
+                "mode": cfg.logger.mode,
+                "config": dict(cfg),
+                "project": cfg.logger.project_name,
+                "group": cfg.logger.group_name,
+            },
+        )
+
+    # Set seeds
+    torch.manual_seed(cfg.env.seed)
+    np.random.seed(cfg.env.seed)
+    device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
+
+    # Create environments
+    train_env, eval_env = make_environment(
+        cfg,
+        cfg.env.train_num_envs,
+        cfg.env.eval_num_envs,
+        logger=logger,
+    )
+
+    # Create replay buffer
+    replay_buffer = make_replay_buffer(
+        batch_size=cfg.optim.batch_size,
+        prb=cfg.replay_buffer.prb,
+        buffer_size=cfg.replay_buffer.size,
+        device="cpu",
+    )
+
+    # Create model
+    model = make_discrete_iql_model(cfg, train_env, eval_env, device)
+
+    compile_mode = None
+    if cfg.compile.compile:
+        compile_mode = cfg.compile.compile_mode
+        if compile_mode in ("", None):
+            if cfg.compile.cudagraphs:
+                compile_mode = "default"
+            else:
+                compile_mode = "reduce-overhead"
+
+    # Create collector
+    collector = make_collector(
+        cfg, train_env, actor_model_explore=model[0], compile_mode=compile_mode
+    )
+
+    # Create loss
+    loss_module, target_net_updater = make_discrete_loss(cfg.loss, model, device=device)
+
+    # Create optimizer
+    optimizer_actor, optimizer_critic, optimizer_value = make_iql_optimizer(
+        cfg.optim, loss_module
+    )
+    optimizer = group_optimizers(optimizer_actor, optimizer_critic, optimizer_value)
+    del optimizer_actor, optimizer_critic, optimizer_value
+
+    def update(sampled_tensordict):
+        optimizer.zero_grad(set_to_none=True)
+        # compute losses
+        actor_loss, _ = loss_module.actor_loss(sampled_tensordict)
+        value_loss, _ = loss_module.value_loss(sampled_tensordict)
+        q_loss, metadata = loss_module.qvalue_loss(sampled_tensordict)
+        (actor_loss + value_loss + q_loss).backward()
+        optimizer.step()
+
+        # update qnet_target params
+        target_net_updater.step()
+        metadata.update(
+            {"actor_loss": actor_loss, "value_loss": value_loss, "q_loss": q_loss}
+        )
+        return TensorDict(metadata).detach()
+
+    if cfg.compile.compile:
+        update = torch.compile(update, mode=compile_mode)
+    if cfg.compile.cudagraphs:
+        warnings.warn(
+            "CudaGraphModule is experimental and may lead to silently wrong results. Use with caution.",
+            category=UserWarning,
+        )
+        update = CudaGraphModule(update, warmup=50)
+
+    # Main loop
+    collected_frames = 0
+    pbar = tqdm.tqdm(total=cfg.collector.total_frames)
+
+    init_random_frames = cfg.collector.init_random_frames
+    num_updates = int(cfg.collector.frames_per_batch * cfg.optim.utd_ratio)
+    prb = cfg.replay_buffer.prb
+    eval_iter = cfg.logger.eval_iter
+    frames_per_batch = cfg.collector.frames_per_batch
+    eval_rollout_steps = cfg.collector.max_frames_per_traj
+
+    collector_iter = iter(collector)
+    total_iter = len(collector)
+    for _ in range(total_iter):
+        timeit.printevery(1000, total_iter, erase=True)
+
+        with timeit("collection"):
+            tensordict = next(collector_iter)
+        current_frames = tensordict.numel()
+        pbar.update(current_frames)
+
+        # update weights of the inference policy
+        collector.update_policy_weights_()
+
+        with timeit("buffer - extend"):
+            tensordict = tensordict.reshape(-1)
+
+            # add to replay buffer
+            replay_buffer.extend(tensordict)
+        collected_frames += current_frames
+
+        # optimization steps
+        with timeit("training"):
+            if collected_frames >= init_random_frames:
+                for _ in range(num_updates):
+                    # sample from replay buffer
+                    with timeit("buffer - sample"):
+                        sampled_tensordict = replay_buffer.sample().to(device)
+
+                    with timeit("training - update"):
+                        torch.compiler.cudagraph_mark_step_begin()
+                        metadata = update(sampled_tensordict)
+                    # update priority
+                    if prb:
+                        sampled_tensordict.set(
+                            loss_module.tensor_keys.priority,
+                            metadata.pop("td_error").detach().max(0).values,
+                        )
+                        replay_buffer.update_priority(sampled_tensordict)
+
+        episode_rewards = tensordict["next", "episode_reward"][
+            tensordict["next", "done"]
+        ]
+
+        metrics_to_log = {}
+        # Evaluation
+        if abs(collected_frames % eval_iter) < frames_per_batch:
+            with set_exploration_type(
+                ExplorationType.DETERMINISTIC
+            ), torch.no_grad(), timeit("eval"):
+                eval_rollout = eval_env.rollout(
+                    eval_rollout_steps,
+                    model[0],
+                    auto_cast_to_device=True,
+                    break_when_any_done=True,
+                )
+                eval_env.apply(dump_video)
+                eval_reward = eval_rollout["next", "reward"].sum(-2).mean().item()
+                metrics_to_log["eval/reward"] = eval_reward
+
+        # Logging
+        if len(episode_rewards) > 0:
+            episode_length = tensordict["next", "step_count"][
+                tensordict["next", "done"]
+            ]
+            metrics_to_log["train/reward"] = episode_rewards.mean().item()
+            metrics_to_log["train/episode_length"] = episode_length.sum().item() / len(
+                episode_length
+            )
+        if collected_frames >= init_random_frames:
+            metrics_to_log["train/q_loss"] = metadata["q_loss"]
+            metrics_to_log["train/actor_loss"] = metadata["actor_loss"]
+            metrics_to_log["train/value_loss"] = metadata["value_loss"]
+        if logger is not None:
+            metrics_to_log.update(timeit.todict(prefix="time"))
+            metrics_to_log["time/speed"] = pbar.format_dict["rate"]
+            log_metrics(logger, metrics_to_log, collected_frames)
+
+    collector.shutdown()
+
+
+if __name__ == "__main__":
+    main()