add offline iql example

pytorch · vmoens · Dec 14, 2023 · Mar 28, 2023 · Mar 28, 2023 · Mar 28, 2023
commit bcf6d46ea553415d60f2e26f531590be1daa6a1a
diff --git a/examples/iql/iql_offline.py b/examples/iql/iql_offline.py
@@ -0,0 +1,92 @@
+# 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 offline IQL training script.
+
+The helper functions are coded in the utils.py associated with this script.
+
+"""
+
+import hydra
+import torch
+import tqdm
+from torchrl.envs.utils import set_exploration_mode
+
+from utils import (
+    get_stats,
+    make_iql_model,
+    make_iql_optimizer,
+    make_logger,
+    make_loss,
+    make_offline_replay_buffer,
+    make_parallel_env,
+)
+
+
+@hydra.main(config_path=".", config_name="offline_config")
+def main(cfg: "DictConfig"):  # noqa: F821
+
+    model_device = cfg.optim.device
+
+    state_dict = get_stats(cfg.env)
+    evaluation_env = make_parallel_env(cfg.env, state_dict=state_dict)
+    logger = make_logger(cfg.logger)
+    replay_buffer = make_offline_replay_buffer(cfg.replay_buffer, state_dict)
+
+    actor_network, qvalue_network, value_network = make_iql_model(cfg)
+    policy = actor_network.to(model_device)
+    qvalue_network = qvalue_network.to(model_device)
+    value_network = value_network.to(model_device)
+
+    loss, target_net_updater = make_loss(
+        cfg.loss, policy, qvalue_network, value_network
+    )
+    optim = make_iql_optimizer(cfg.optim, policy, qvalue_network, value_network)
+
+    pbar = tqdm.tqdm(total=cfg.optim.gradient_steps)
+
+    r0 = None
+    l0 = None
+
+    for i in range(cfg.optim.gradient_steps):
+        pbar.update(i)
+        data = replay_buffer.sample()
+        # loss
+        loss_vals = loss(data)
+        # backprop
+        actor_loss = loss_vals["loss_actor"]
+        q_loss = loss_vals["loss_qvalue"]
+        value_loss = loss_vals["loss_value"]
+        loss_val = actor_loss + q_loss + value_loss
+
+        optim.zero_grad()
+        loss_val.backward()
+        optim.step()
+        target_net_updater.step()
+
+        # evaluation
+        if i % cfg.env.evaluation_interval == 0:
+            with set_exploration_mode("random"), torch.no_grad():
+                eval_td = evaluation_env.rollout(
+                    max_steps=1000, policy=policy, auto_cast_to_device=True
+                )
+
+        if r0 is None:
+            r0 = eval_td["reward"].mean().item()
+        if l0 is None:
+            l0 = loss_val.item()
+
+        for key, value in loss_vals.items():
+            logger.log_scalar(key, value.item(), i)
+        logger.log_scalar("reward_evaluation", eval_td["reward"].mean().item(), i)
+
+        pbar.set_description(
+            f"loss: {loss_val.item(): 4.4f} (init: {l0: 4.4f}), reward: {eval_td['reward'].mean(): 4.4f} (init={r0: 4.4f})"
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/iql/iql_online.py b/examples/iql/iql_online.py
@@ -4,7 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 """IQL Example.
 
-This is a self-contained example of a IQL training script.
+This is a self-contained example of an online IQL training script.
 
 It works across Gym and DM-control over a variety of tasks.
 

diff --git a/examples/iql/offline_config.yaml b/examples/iql/offline_config.yaml
@@ -0,0 +1,55 @@
+# Task and env
+env:
+  env_name: Hopper-v3
+  env_task: ""
+  env_library: gym
+  record_video: 0
+  n_samples_stats: 1000
+  frame_skip: 1
+  from_pixels: False
+  num_eval_envs: 1
+  reward_scaling:
+  noop: 1
+  seed: 0
+  evaluation_interval: 1000
+
+# Eval
+recorder:
+  video: False
+  interval: 10000 # record interval in frames
+  frames: 10000
+
+# logger
+logger:
+  backend: wandb
+  exp_name: iql_hopper-medium-v2
+
+# Buffer
+replay_buffer:
+  dataset: hopper-medium-v2
+  batch_size: 256
+
+# Optimization
+optim:
+  device: cpu
+  lr: 3e-4
+  weight_decay: 0.0
+  batch_size: 256
+  lr_scheduler: ""
+  gradient_steps: 1000000
+
+
+# Policy and model
+model:
+  activation: relu
+  default_policy_scale: 1.0
+  scale_lb: 0.1
+
+# loss
+loss: 
+  loss_function: smooth_l1
+  gamma: 0.99
+  tau: 0.05
+  # IQL hyperparameter
+  temperature: 3.0
+  expectile: 0.7
diff --git a/examples/iql/online_config.yaml b/examples/iql/online_config.yaml
@@ -51,7 +51,6 @@ optim:
   batch_size: 256
   lr_scheduler: ""
   optim_steps_per_batch: 1000
-  policy_update_delay: 2
 
 # Policy and model
 model:

diff --git a/examples/iql/utils.py b/examples/iql/utils.py
@@ -12,6 +12,8 @@
     MultiStep,
     TensorDictReplayBuffer,
 )
+from torchrl.data.datasets.d4rl import D4RLExperienceReplay
+from torchrl.data.replay_buffers import SamplerWithoutReplacement
 from torchrl.data.replay_buffers.samplers import PrioritizedSampler, RandomSampler
 from torchrl.envs import (
     CatFrames,
@@ -22,6 +24,7 @@
     NoopResetEnv,
     ObservationNorm,
     ParallelEnv,
+    RenameTransform,
     Resize,
     RewardScaling,
     ToTensorImage,
@@ -169,7 +172,7 @@ def make_transformed_env_states(base_env, env_cfg):
 
 
 def make_parallel_env(env_cfg, state_dict):
-    num_envs = env_cfg.num_envs
+    num_envs = env_cfg.num_eval_envs
     env = make_transformed_env(
         ParallelEnv(num_envs, EnvCreator(lambda: make_base_env(env_cfg))), env_cfg
     )
@@ -242,6 +245,45 @@ def make_replay_buffer(rb_cfg):
     )
 
 
+def make_offline_replay_buffer(rb_cfg, state_dict):
+
+    data = D4RLExperienceReplay(
+        rb_cfg.dataset,
+        split_trajs=False,
+        batch_size=rb_cfg.batch_size,
+        sampler=SamplerWithoutReplacement(drop_last=False),
+    )
+    data.append_transform(
+        RewardScaling(
+            loc=state_dict["transforms.0.loc"],
+            scale=state_dict["transforms.0.scale"],
+            standard_normal=state_dict["transforms.0.standard_normal"],
+        )
+    )
+    data.append_transform(
+        RenameTransform(
+            ["observation", ("next", "observation")],
+            ["observation_vector", ("next", "observation_vector")],
+        )
+    )
+    data.append_transform(
+        ObservationNorm(
+            in_keys=["observation_vector", ("next", "observation_vector")],
+            loc=state_dict["transforms.2.loc"],
+            scale=state_dict["transforms.2.scale"],
+            standard_normal=state_dict["transforms.2.standard_normal"],
+        )
+    )
+    data.append_transform(
+        DoubleToFloat(
+            in_keys=["observation_vector", ("next", "observation_vector")],
+            in_keys_inv=[],
+        )
+    )
+
+    return data
+
+
 # ====================================================================
 # Model
 # -----