torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

sota-implementations/decision_transformer/utils.py

@@ -0,0 +1,562 @@
+# 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 os
+from pathlib import Path
+
+import torch.nn
+
+import torch.optim
+from lamb import Lamb
+from tensordict.nn import TensorDictModule
+
+from torchrl.collectors import SyncDataCollector
+from torchrl.data import (
+    LazyMemmapStorage,
+    RoundRobinWriter,
+    TensorDictReplayBuffer,
+    TensorStorage,
+)
+from torchrl.data.datasets.d4rl import D4RLExperienceReplay
+from torchrl.data.replay_buffers import RandomSampler
+from torchrl.envs import (
+    CatFrames,
+    Compose,
+    DoubleToFloat,
+    EnvCreator,
+    ExcludeTransform,
+    ObservationNorm,
+    ParallelEnv,
+    RandomCropTensorDict,
+    RenameTransform,
+    Reward2GoTransform,
+    RewardScaling,
+    RewardSum,
+    TargetReturn,
+    TensorDictPrimer,
+    TransformedEnv,
+    UnsqueezeTransform,
+)
+from torchrl.envs.libs.dm_control import DMControlEnv
+from torchrl.envs.libs.gym import set_gym_backend
+from torchrl.envs.utils import ExplorationType, set_exploration_type
+from torchrl.modules import (
+    DTActor,
+    OnlineDTActor,
+    ProbabilisticActor,
+    TanhDelta,
+    TanhNormal,
+)
+
+from torchrl.objectives import DTLoss, OnlineDTLoss
+from torchrl.record import VideoRecorder
+from torchrl.record.loggers import generate_exp_name, get_logger
+from torchrl.trainers.helpers.envs import LIBS
+
+# ====================================================================
+# Environment utils
+# -----------------
+
+
+def make_base_env(env_cfg, from_pixels=False, device=None):
+    set_gym_backend(env_cfg.backend).set()
+
+    env_library = LIBS[env_cfg.library]
+    env_name = env_cfg.name
+    frame_skip = env_cfg.frame_skip
+
+    env_kwargs = {
+        "env_name": env_name,
+        "frame_skip": frame_skip,
+        "from_pixels": from_pixels,
+        "pixels_only": False,
+    }
+    if env_library is DMControlEnv:
+        env_task = env_cfg.task
+        env_kwargs.update({"task_name": env_task})
+    env = env_library(**env_kwargs, device=device)
+    return env
+
+
+def make_transformed_env(base_env, env_cfg, obs_loc, obs_std, train=False):
+    transformed_env = TransformedEnv(base_env)
+    transformed_env.append_transform(
+        RewardScaling(
+            loc=0,
+            scale=env_cfg.reward_scaling,
+            in_keys=["reward"],
+            standard_normal=False,
+        )
+    )
+    if train:
+        transformed_env.append_transform(
+            TargetReturn(
+                env_cfg.collect_target_return * env_cfg.reward_scaling,
+                out_keys=["return_to_go"],
+                mode=env_cfg.target_return_mode,
+            )
+        )
+    else:
+        transformed_env.append_transform(
+            TargetReturn(
+                env_cfg.eval_target_return * env_cfg.reward_scaling,
+                out_keys=["return_to_go"],
+                mode=env_cfg.target_return_mode,
+            )
+        )
+
+    # copy action from the input tensordict to the output
+    transformed_env.append_transform(TensorDictPrimer(base_env.full_action_spec))
+
+    transformed_env.append_transform(DoubleToFloat())
+    obsnorm = ObservationNorm(
+        loc=obs_loc, scale=obs_std, in_keys="observation", standard_normal=True
+    )
+    transformed_env.append_transform(obsnorm)
+    transformed_env.append_transform(
+        UnsqueezeTransform(
+            -2,
+            in_keys=["observation", "action", "return_to_go"],
+            out_keys=["observation_cat", "action_cat", "return_to_go_cat"],
+        )
+    )
+    transformed_env.append_transform(
+        CatFrames(
+            in_keys=["observation_cat", "action_cat", "return_to_go_cat"],
+            N=env_cfg.stacked_frames,
+            dim=-2,
+            padding="constant",
+        )
+    )
+
+    if train:
+        transformed_env.append_transform(RewardSum())
+
+    return transformed_env
+
+
+def make_parallel_env(
+    env_cfg, obs_loc, obs_std, train=False, from_pixels=False, device=None
+):
+    if train:
+        num_envs = env_cfg.num_train_envs
+    else:
+        num_envs = env_cfg.num_eval_envs
+
+    def make_env():
+        with set_gym_backend(env_cfg.backend):
+            return make_base_env(env_cfg, from_pixels=from_pixels, device="cpu")
+
+    env = make_transformed_env(
+        ParallelEnv(
+            num_envs, EnvCreator(make_env), serial_for_single=True, device=device
+        ),
+        env_cfg,
+        obs_loc,
+        obs_std,
+        train,
+    )
+    env.start()
+    return env
+
+
+def make_env(env_cfg, obs_loc, obs_std, train=False, from_pixels=False, device=None):
+    return make_parallel_env(
+        env_cfg,
+        obs_loc,
+        obs_std,
+        train=train,
+        from_pixels=from_pixels,
+        device=device,
+    )
+
+
+# ====================================================================
+# Collector and replay buffer
+# ---------------------------
+
+
+def make_collector(cfg, policy):
+    exclude_target_return = ExcludeTransform(
+        "return_to_go",
+        ("next", "return_to_go"),
+        ("next", "action"),
+        ("next", "observation"),
+        "scale",
+        "loc",
+    )
+    cat = CatFrames(
+        in_keys=["action"], out_keys=["action_cat"], N=20, dim=-2, padding="constant"
+    )
+    transforms = Compose(
+        exclude_target_return,
+        cat,
+    )
+    collector_cfg = cfg.collector
+    collector_class = SyncDataCollector
+    collector = collector_class(
+        make_env(cfg.env, train=True),
+        policy,
+        frames_per_batch=collector_cfg.frames_per_batch,
+        total_frames=collector_cfg.total_frames,
+        device=collector_cfg.devices,
+        max_frames_per_traj=collector_cfg.max_frames_per_traj,
+        postproc=transforms,
+    )
+    return collector
+
+
+def make_offline_replay_buffer(rb_cfg, reward_scaling):
+    r2g = Reward2GoTransform(
+        gamma=1.0,
+        in_keys=[("next", "reward"), "reward"],
+        out_keys=[("next", "return_to_go"), "return_to_go"],
+    )
+    reward_scale = RewardScaling(
+        loc=0,
+        scale=reward_scaling,
+        in_keys=[("next", "return_to_go"), "return_to_go"],
+        standard_normal=False,
+    )
+    crop_seq = RandomCropTensorDict(sub_seq_len=rb_cfg.stacked_frames, sample_dim=-1)
+    d2f = DoubleToFloat()
+    rename = RenameTransform(
+        in_keys=[
+            "action",
+            "observation",
+            "return_to_go",
+            ("next", "return_to_go"),
+            ("next", "observation"),
+        ],
+        out_keys=[
+            "action_cat",
+            "observation_cat",
+            "return_to_go_cat",
+            ("next", "return_to_go_cat"),
+            ("next", "observation_cat"),
+        ],
+    )
+    exclude = ExcludeTransform(
+        "terminal",
+        "info",
+        ("next", "timeout"),
+        ("next", "terminal"),
+        ("next", "observation"),
+        ("next", "info"),
+    )
+
+    transforms = Compose(
+        r2g,
+        crop_seq,
+        reward_scale,
+        d2f,
+        rename,
+        exclude,
+    )
+    data = D4RLExperienceReplay(
+        dataset_id=rb_cfg.dataset,
+        split_trajs=True,
+        batch_size=rb_cfg.batch_size,
+        sampler=RandomSampler(),  # SamplerWithoutReplacement(drop_last=False),
+        transform=None,
+        use_truncated_as_done=True,
+        direct_download=True,
+        prefetch=4,
+        writer=RoundRobinWriter(),
+        root=Path(os.environ["HOME"]) / ".cache" / "torchrl" / "data" / "d4rl",
+    )
+
+    # since we're not extending the data, adding keys can only be done via
+    # the creation of a new storage
+    data_memmap = data[:]
+    with data_memmap.unlock_():
+        data_memmap = r2g.inv(data_memmap)
+        data._storage = TensorStorage(data_memmap)
+
+    loc = data[:]["observation"].flatten(0, -2).mean(axis=0).float()
+    std = data[:]["observation"].flatten(0, -2).std(axis=0).float()
+
+    obsnorm = ObservationNorm(
+        loc=loc,
+        scale=std,
+        in_keys=["observation_cat", ("next", "observation_cat")],
+        standard_normal=True,
+    )
+    for t in transforms:
+        data.append_transform(t)
+    data.append_transform(obsnorm)
+    return data, loc, std
+
+
+def make_online_replay_buffer(offline_buffer, rb_cfg, reward_scaling=0.001):
+    r2g = Reward2GoTransform(gamma=1.0, out_keys=["return_to_go"])
+    reward_scale = RewardScaling(
+        loc=0,
+        scale=reward_scaling,
+        in_keys=["return_to_go"],
+        out_keys=["return_to_go"],
+        standard_normal=False,
+    )
+    catframes = CatFrames(
+        in_keys=["return_to_go"],
+        out_keys=["return_to_go_cat"],
+        N=rb_cfg.stacked_frames,
+        dim=-2,
+        padding="constant",
+        as_inverse=True,
+    )
+    transforms = Compose(
+        r2g,
+        reward_scale,
+        catframes,
+    )
+    storage = LazyMemmapStorage(
+        max_size=rb_cfg.capacity,
+        scratch_dir=rb_cfg.scratch_dir,
+        device=rb_cfg.device,
+    )
+
+    replay_buffer = TensorDictReplayBuffer(
+        pin_memory=False,
+        prefetch=rb_cfg.prefetch,
+        storage=storage,
+        batch_size=rb_cfg.batch_size,
+    )
+    # init buffer with offline data
+    offline_data = offline_buffer[:100000]
+    offline_data.del_("index")
+    replay_buffer.extend(offline_data.clone().detach().to_tensordict())
+    # add transforms after offline data extension to not trigger reward-to-go calculation
+    replay_buffer.append_transform(transforms)
+
+    return replay_buffer
+
+
+# ====================================================================
+# Model
+# -----
+
+
+def make_odt_model(cfg, device: torch.device | None = None) -> TensorDictModule:
+    env_cfg = cfg.env
+    proof_environment = make_transformed_env(
+        make_base_env(env_cfg), env_cfg, obs_loc=0, obs_std=1
+    )
+
+    action_spec = proof_environment.action_spec_unbatched
+    for key, value in proof_environment.observation_spec_unbatched.items():
+        if key == "observation":
+            state_dim = value.shape[-1]
+    in_keys = [
+        "observation_cat",
+        "action_cat",
+        "return_to_go_cat",
+    ]
+
+    actor_net = OnlineDTActor(
+        state_dim=state_dim,
+        action_dim=action_spec.shape[-1],
+        transformer_config=cfg.transformer,
+        device=device,
+    )
+
+    actor_module = TensorDictModule(
+        actor_net,
+        in_keys=in_keys,
+        out_keys=[
+            "loc",
+            "scale",
+        ],
+    )
+    dist_class = TanhNormal
+    dist_kwargs = {
+        "low": -torch.ones((), device=device),
+        "high": torch.ones((), device=device),
+        "tanh_loc": False,
+        "upscale": torch.full((), 5, device=device),
+        # "safe_tanh": not cfg.compile.compile,
+    }
+
+    actor = ProbabilisticActor(
+        spec=action_spec,
+        in_keys=["loc", "scale"],
+        out_keys=["action"],
+        module=actor_module,
+        distribution_class=dist_class,
+        distribution_kwargs=dist_kwargs,
+        cache_dist=False,
+        return_log_prob=False,
+    )
+
+    # init the lazy layers
+    with torch.no_grad(), set_exploration_type(ExplorationType.RANDOM):
+        td = proof_environment.rollout(max_steps=100)
+        td["action"] = td["next", "action"]
+        actor(td.to(device))
+
+    return actor
+
+
+def make_dt_model(cfg, device: torch.device | None = None):
+    env_cfg = cfg.env
+    proof_environment = make_transformed_env(
+        make_base_env(env_cfg), env_cfg, obs_loc=0, obs_std=1
+    )
+
+    action_spec = proof_environment.action_spec_unbatched
+    in_keys = [
+        "observation_cat",
+        "action_cat",
+        "return_to_go_cat",
+    ]
+
+    actor_net = DTActor(
+        state_dim=proof_environment.observation_spec_unbatched["observation"].shape[-1],
+        action_dim=action_spec.shape[-1],
+        transformer_config=cfg.transformer,
+        device=device,
+    )
+
+    actor_module = TensorDictModule(
+        actor_net,
+        in_keys=in_keys,
+        out_keys=["param"],
+    )
+    dist_class = TanhDelta
+    dist_kwargs = {
+        "low": action_spec.space.low.to(device),
+        "high": action_spec.space.high.to(device),
+        "safe": not cfg.compile.compile,
+    }
+
+    actor = ProbabilisticActor(
+        spec=action_spec.to(device),
+        in_keys=["param"],
+        out_keys=["action"],
+        module=actor_module,
+        distribution_class=dist_class,
+        distribution_kwargs=dist_kwargs,
+        cache_dist=False,
+        return_log_prob=False,
+    )
+
+    # init the lazy layers
+    with torch.no_grad(), set_exploration_type(ExplorationType.RANDOM):
+        td = proof_environment.fake_tensordict()
+        td = td.expand((100, *td.shape))
+        td["action"] = td["next", "action"]
+        actor(td.to(device))
+
+    return actor
+
+
+# ====================================================================
+# Online Decision Transformer Loss
+# ---------
+
+
+def make_odt_loss(loss_cfg, actor_network):
+    loss = OnlineDTLoss(
+        actor_network,
+        alpha_init=loss_cfg.alpha_init,
+        target_entropy=loss_cfg.target_entropy,
+    )
+    loss.set_keys(action_target="action_cat")
+    return loss
+
+
+def make_dt_loss(loss_cfg, actor_network, device: torch.device | None = None):
+    loss = DTLoss(
+        actor_network,
+        loss_function=loss_cfg.loss_function,
+        device=device,
+    )
+    loss.set_keys(action_target="action_cat")
+    return loss
+
+
+def make_odt_optimizer(optim_cfg, loss_module):
+    if optim_cfg.optimizer == "lamb":
+        dt_optimizer = Lamb(
+            loss_module.actor_network_params.flatten_keys().values(),
+            lr=torch.as_tensor(
+                optim_cfg.lr, device=next(loss_module.parameters()).device
+            ),
+            weight_decay=optim_cfg.weight_decay,
+            eps=1.0e-8,
+        )
+    elif optim_cfg.optimizer == "adam":
+        dt_optimizer = torch.optim.Adam(
+            loss_module.actor_network_params.flatten_keys().values(),
+            lr=torch.as_tensor(
+                optim_cfg.lr, device=next(loss_module.parameters()).device
+            ),
+            weight_decay=optim_cfg.weight_decay,
+            eps=1.0e-8,
+        )
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(
+        dt_optimizer, lambda steps: min((steps + 1) / optim_cfg.warmup_steps, 1)
+    )
+
+    log_temp_optimizer = torch.optim.Adam(
+        [loss_module.log_alpha],
+        lr=torch.as_tensor(1e-4, device=next(loss_module.parameters()).device),
+        betas=[0.9, 0.999],
+    )
+
+    return dt_optimizer, log_temp_optimizer, scheduler
+
+
+def make_dt_optimizer(optim_cfg, loss_module, device):
+    dt_optimizer = torch.optim.Adam(
+        loss_module.actor_network_params.flatten_keys().values(),
+        lr=torch.tensor(optim_cfg.lr, device=device),
+        weight_decay=optim_cfg.weight_decay,
+        eps=1.0e-8,
+    )
+    scheduler = torch.optim.lr_scheduler.LambdaLR(
+        dt_optimizer, lambda steps: min((steps + 1) / optim_cfg.warmup_steps, 1)
+    )
+
+    return dt_optimizer, scheduler
+
+
+# ====================================================================
+# Logging and recording
+# ---------------------
+
+
+def make_logger(cfg):
+    if not cfg.logger.backend:
+        return None
+    exp_name = generate_exp_name(cfg.logger.model_name, cfg.logger.exp_name)
+    logger = get_logger(
+        cfg.logger.backend,
+        logger_name=cfg.logger.model_name,
+        experiment_name=exp_name,
+        wandb_kwargs={
+            "config": dict(cfg),
+            "project": cfg.logger.project_name,
+            "group": cfg.logger.group_name,
+        },
+    )
+    return logger
+
+
+# ====================================================================
+# General utils
+# ---------
+
+
+def log_metrics(logger, metrics, step):
+    for metric_name, metric_value in metrics.items():
+        logger.log_scalar(metric_name, metric_value, step)
+
+
+def dump_video(module):
+    if isinstance(module, VideoRecorder):
+        module.dump()