torchrl 0.11.0__cp314-cp314-win_amd64.whl

sota-implementations/ppo/utils_atari.py

@@ -0,0 +1,238 @@
+# 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.data.tensor_specs import CategoricalBox
+from torchrl.envs import (
+    CatFrames,
+    DoubleToFloat,
+    EndOfLifeTransform,
+    EnvCreator,
+    ExplorationType,
+    GrayScale,
+    GymEnv,
+    NoopResetEnv,
+    ParallelEnv,
+    RenameTransform,
+    Resize,
+    RewardSum,
+    set_gym_backend,
+    SignTransform,
+    StepCounter,
+    ToTensorImage,
+    TransformedEnv,
+    VecNorm,
+)
+from torchrl.modules import (
+    ActorValueOperator,
+    ConvNet,
+    MLP,
+    ProbabilisticActor,
+    TanhNormal,
+    ValueOperator,
+)
+from torchrl.record import VideoRecorder
+
+
+# ====================================================================
+# Environment utils
+# --------------------------------------------------------------------
+
+
+def make_base_env(
+    env_name="BreakoutNoFrameskip-v4",
+    frame_skip=4,
+    gym_backend="gymnasium",
+    is_test=False,
+):
+    with set_gym_backend(gym_backend):
+        env = GymEnv(
+            env_name,
+            frame_skip=frame_skip,
+            from_pixels=True,
+            pixels_only=False,
+            device="cpu",
+            categorical_action_encoding=True,
+        )
+    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)),
+        serial_for_single=True,
+        device=device,
+    )
+    env = TransformedEnv(env)
+    env.append_transform(RenameTransform(in_keys=["pixels"], out_keys=["pixels_int"]))
+    env.append_transform(ToTensorImage(in_keys=["pixels_int"], out_keys=["pixels"]))
+    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(DoubleToFloat())
+    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 = torch.distributions.Categorical
+        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, 1, device=device, 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(10)
+        actor_critic(td)
+        del td
+
+    actor = actor_critic.get_policy_operator()
+    critic = actor_critic.get_value_operator()
+
+    del proof_environment
+
+    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,
+        )
+        test_env.apply(dump_video)
+        reward = td_test["next", "episode_reward"][td_test["next", "done"]]
+        test_rewards.append(reward.cpu())
+    del td_test
+    return torch.cat(test_rewards, 0).mean()

sota-implementations/ppo/utils_mujoco.py

@@ -0,0 +1,152 @@
+# 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 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: bool = False):
+    env = GymEnv(env_name, device=device, from_pixels=from_pixels, pixels_only=False)
+    env = TransformedEnv(env)
+    env.append_transform(VecNorm(in_keys=["observation"], decay=0.99999, eps=1e-2))
+    env.append_transform(ClipTransform(in_keys=["observation"], low=-10, high=10))
+    env.append_transform(RewardSum())
+    env.append_transform(StepCounter())
+    env.append_transform(DoubleToFloat(in_keys=["observation"]))
+    return env
+
+
+# ====================================================================
+# Model utils
+# --------------------------------------------------------------------
+
+
+def make_ppo_models_state(proof_environment, device):
+
+    # 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,
+    }
+
+    # 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], scale_lb=1e-8
+        ).to(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):
+    proof_environment = make_env(env_name, device=device)
+    actor, critic = make_ppo_models_state(proof_environment, device=device)
+    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.append(reward.cpu())
+        test_env.apply(dump_video)
+    del td_test
+    return torch.cat(test_rewards, 0).mean()

sota-implementations/ppo_trainer/train.py

@@ -0,0 +1,21 @@
+# 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.
+
+import hydra
+import torchrl
+from torchrl.trainers.algorithms.configs import *  # noqa: F401, F403
+
+
+@hydra.main(config_path="config", config_name="config", version_base="1.1")
+def main(cfg):
+    def print_reward(td):
+        torchrl.logger.info(f"reward: {td['next', 'reward'].mean(): 4.4f}")
+
+    trainer = hydra.utils.instantiate(cfg.trainer)
+    trainer.register_op(dest="batch_process", op=print_reward)
+    trainer.train()
+
+
+if __name__ == "__main__":
+    main()

sota-implementations/redq/README.md

@@ -0,0 +1,7 @@
+# REDQ 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/redq/redq.py

@@ -0,0 +1,199 @@
+# 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 uuid
+from datetime import datetime
+
+import hydra
+import torch.cuda
+from tensordict.nn import TensorDictSequential
+from torchrl.envs import EnvCreator, ParallelEnv
+from torchrl.envs.transforms import RewardScaling, TransformedEnv
+from torchrl.envs.utils import ExplorationType, set_exploration_type
+from torchrl.modules import OrnsteinUhlenbeckProcessModule
+from torchrl.record import VideoRecorder
+from torchrl.record.loggers import get_logger
+from utils import (
+    correct_for_frame_skip,
+    get_norm_state_dict,
+    initialize_observation_norm_transforms,
+    make_collector_offpolicy,
+    make_redq_loss,
+    make_redq_model,
+    make_replay_buffer,
+    make_trainer,
+    parallel_env_constructor,
+    retrieve_observation_norms_state_dict,
+    transformed_env_constructor,
+)
+
+DEFAULT_REWARD_SCALING = {
+    "Hopper-v1": 5,
+    "Walker2d-v1": 5,
+    "HalfCheetah-v1": 5,
+    "cheetah": 5,
+    "Ant-v2": 5,
+    "Humanoid-v2": 20,
+    "humanoid": 100,
+}
+
+
+@hydra.main(version_base="1.1", config_path="", config_name="config")
+def main(cfg: DictConfig):  # noqa: F821
+
+    cfg = correct_for_frame_skip(cfg)
+
+    if not isinstance(cfg.env.reward_scaling, float):
+        cfg.env.reward_scaling = DEFAULT_REWARD_SCALING.get(cfg.env.name, 5.0)
+        cfg.env.reward_loc = 0.0
+
+    device = (
+        torch.device("cpu")
+        if torch.cuda.device_count() == 0
+        else torch.device("cuda:0")
+    )
+
+    exp_name = "_".join(
+        [
+            "REDQ",
+            cfg.logger.exp_name,
+            str(uuid.uuid4())[:8],
+            datetime.now().strftime("%y_%m_%d-%H_%M_%S"),
+        ]
+    )
+
+    if cfg.logger.backend:
+        logger = get_logger(
+            logger_type=cfg.logger.backend,
+            logger_name="redq_logging",
+            experiment_name=exp_name,
+            wandb_kwargs={
+                "mode": cfg.logger.mode,
+                "config": dict(cfg),
+                "project": cfg.logger.project_name,
+                "group": cfg.logger.group_name,
+            },
+        )
+    else:
+        logger = None
+
+    key, init_env_steps, stats = None, None, None
+    if not cfg.env.vecnorm and cfg.env.norm_stats:
+        key = (
+            ("next", "pixels")
+            if cfg.env.from_pixels
+            else ("next", "observation_vector")
+        )
+        init_env_steps = cfg.env.init_env_steps
+        stats = {"loc": None, "scale": None}
+    elif cfg.env.from_pixels:
+        stats = {"loc": 0.5, "scale": 0.5}
+
+    proof_env = transformed_env_constructor(
+        cfg=cfg,
+        use_env_creator=False,
+        stats=stats,
+    )()
+    initialize_observation_norm_transforms(
+        proof_environment=proof_env, num_iter=init_env_steps, key=key
+    )
+    _, obs_norm_state_dict = retrieve_observation_norms_state_dict(proof_env)[0]
+
+    model = make_redq_model(
+        proof_env,
+        cfg=cfg,
+        device=device,
+    )
+    loss_module, target_net_updater = make_redq_loss(model, cfg)
+
+    actor_model_explore = model[0]
+    if cfg.exploration.ou_exploration:
+        if cfg.exploration.gSDE:
+            raise RuntimeError("gSDE and ou_exploration are incompatible")
+        actor_model_explore = TensorDictSequential(
+            actor_model_explore,
+            OrnsteinUhlenbeckProcessModule(
+                spec=actor_model_explore.spec,
+                annealing_num_steps=cfg.exploration.annealing_frames,
+                sigma=cfg.exploration.ou_sigma,
+                theta=cfg.exploration.ou_theta,
+                device=device,
+            ),
+        )
+    if device == torch.device("cpu"):
+        # mostly for debugging
+        actor_model_explore.share_memory()
+
+    if cfg.exploration.gSDE:
+        with torch.no_grad(), set_exploration_type(ExplorationType.RANDOM):
+            # get dimensions to build the parallel env
+            proof_td = actor_model_explore(proof_env.reset().to(device))
+        action_dim_gsde, state_dim_gsde = proof_td.get("_eps_gSDE").shape[-2:]
+        del proof_td
+    else:
+        action_dim_gsde, state_dim_gsde = None, None
+
+    proof_env.close()
+    create_env_fn = parallel_env_constructor(
+        cfg=cfg,
+        obs_norm_state_dict=obs_norm_state_dict,
+        action_dim_gsde=action_dim_gsde,
+        state_dim_gsde=state_dim_gsde,
+    )
+
+    collector = make_collector_offpolicy(
+        make_env=create_env_fn,
+        actor_model_explore=actor_model_explore,
+        cfg=cfg,
+    )
+
+    replay_buffer = make_replay_buffer("cpu", cfg)
+
+    recorder = transformed_env_constructor(
+        cfg,
+        video_tag="rendering/test",
+        norm_obs_only=True,
+        obs_norm_state_dict=obs_norm_state_dict,
+        logger=logger,
+        use_env_creator=False,
+    )()
+    if isinstance(create_env_fn, ParallelEnv):
+        raise NotImplementedError("This behavior is deprecated")
+    elif isinstance(create_env_fn, EnvCreator):
+        recorder.transform[1:].load_state_dict(
+            get_norm_state_dict(create_env_fn()), strict=False
+        )
+    elif isinstance(create_env_fn, TransformedEnv):
+        recorder.transform = create_env_fn.transform.clone()
+    else:
+        raise NotImplementedError(f"Unsupported env type {type(create_env_fn)}")
+    if logger is not None and cfg.logger.video:
+        recorder.insert_transform(0, VideoRecorder(logger=logger, tag="rendering/test"))
+
+    # reset reward scaling
+    for t in recorder.transform:
+        if isinstance(t, RewardScaling):
+            t.scale.fill_(1.0)
+            t.loc.fill_(0.0)
+
+    trainer = make_trainer(
+        collector=collector,
+        loss_module=loss_module,
+        recorder=recorder,
+        target_net_updater=target_net_updater,
+        policy_exploration=actor_model_explore,
+        replay_buffer=replay_buffer,
+        logger=logger,
+        cfg=cfg,
+    )
+
+    trainer.train()
+    if logger is not None:
+        return (logger.log_dir, trainer._log_dict)
+
+
+if __name__ == "__main__":
+    main()