Heterogenous Pre-trained Transformers (HPT)

To initialize a policy and its environment-specific components, follow these steps:

1. Load a Pretrained Policy: Use Policy.from_pretrained with a Hugging Face path (e.g., hf://liruiw/hpt-base).
2. Initialize Configuration: Use hydra.initialize and hydra.compose to load configurations from the experiments/configs directory. You can override the env parameter to specify the domain (e.g., mujoco_metaworld).
3. Instantiate Dataset: Use hydra.utils.instantiate on the cfg.dataset object, providing the dataset_name, env_rollout_fn, and unpacking **cfg.dataset.
4. Configure Model Architecture:
   • Use utils.update_network_dim(cfg, dataset, policy) to align dimensions.
   • Call policy.init_domain_stem(domain, cfg.stem) to initialize the stem.
   • Call policy.init_domain_head(domain, normalizer, cfg.head) to initialize the head using the dataset's normalizer.
   • Call policy.finalize_modules() to complete the setup.
5. Move to Device: Use policy.to(device) to move the model to a GPU or CPU.

from hydra import compose, initialize
from hpt.models.policy import Policy
from hpt.utils import utils

# 1. Load policy
policy = Policy.from_pretrained("hf://liruiw/hpt-base")
device = 'cuda' if torch.cuda.is_available() else 'cpu'
domain = "mujoco_metaworld"

# 2. Setup Config
with initialize(version_base="1.2", config_path="experiments/configs"):
    cfg = compose(config_name="config", overrides=[f"env={domain}"])

# 3. Setup Dataset
cfg.dataset.episode_cnt = 10
dataset = hydra.utils.instantiate(
    cfg.dataset, 
    dataset_name=domain, 
    env_rollout_fn=cfg.dataset_generator_func, 
    **cfg.dataset
)
normalizer = dataset.get_normalizer()

# 4. Setup Model
utils.update_network_dim(cfg, dataset, policy)
policy.init_domain_stem(domain, cfg.stem)
policy.init_domain_head(domain, normalizer, cfg.head)
policy.finalize_modules()
policy.to(device)

To run a policy in a simulation environment (rollout):

1. Initialize Environment: Instantiate the environment from the ALL_ENVS registry.
2. Prepare Observation: Create an OrderedDict containing the required observation keys (e.g., state and image).
3. Reset Policy: Call policy.reset() before starting the loop.
4. Step Loop: In each step of env.max_path_length:
   • Get action: a = policy.get_action(step_data).
   • Step environment: o, r, done, info = env.step(a).
   • Update observation: Render the new image and update step_data.

from collections import OrderedDict
from env.mujoco.metaworld.envs.mujoco.sawyer_xyz.test_scripted_policies import ALL_ENVS

# Setup
env = ALL_ENVS["reach-v2"]()
env._partially_observable = False
env._freeze_rand_vec = False
env._set_task_called = True

# Initial observation
img = env.sim.render(128, 128, mode="offscreen", camera_name="view_1")[:, :, ::-1].copy()
o = env.reset()
step_data = OrderedDict({"state": o, "image": img})

policy.reset()

# Rollout loop
for _ in range(env.max_path_length):
    a = policy.get_action(step_data)
    o, r, done, info = env.step(a)
    img = env.sim.render(128, 128, mode="offscreen", camera_name="view_1")[:, :, ::-1]
    step_data = OrderedDict({"state": o, "image": img})
    if done:
        break

The training process is controlled via a Hydra configuration file (defaulting to config.yaml in ../experiments/configs). Key configuration sections include:

• train: Controls pretraining paths (pretrained_dir), total iterations (total_iters), and whether to freeze the trunk (freeze_trunk).
• network: Defines the model architecture and whether to finetune encoders (finetune_encoder).
• dataset: Configuration for the dataset instantiation, including dataset_name and env_rollout_fn.
• dataloader / val_dataloader: Standard PyTorch DataLoader parameters.
• optimizer / optimizer_misc: Parameters for the optimizer.
• lr_scheduler: Parameters for the learning rate scheduler.
• warmup_lr: Configuration for the linear warmup steps (step).
• domains: A comma-separated string of domains used to select the active domain.
• output_dir: The directory where logs and the final model.pth are saved.

The evaluation script relies on a Hydra configuration hierarchy. Key configuration keys required by run_eval.py include:

• seed: Integer used for reproducibility.
• output_dir: Directory where results are saved (appended with the seed).
• domains: A comma-separated string of domain names (the first domain in the list is used for initialization).
• network: A configuration object that can be instantiated via hydra.utils.instantiate. It must support:
  • finetune_encoder: Boolean flag.
  • stem: Configuration for the domain stem.
  • head: Configuration for the domain head.
• dataset.image_encoder: Configuration for the image encoder.
• train.pretrained_dir: The directory containing the model.pth file.

Example configuration structure (conceptual):

seed: 42
output_dir: "./results"
domains: "domain1,domain2"
network:
  _target_: hpt.networks.YourPolicyClass
  finetune_encoder: true
  stem: { ... }
  head: { ... }
train:
  pretrained_dir: "/path/to/checkpoints"
dataset:
  image_encoder: { ... }