| import os |
| from typing import List |
| import torch |
| import numpy as np |
|
|
| from omegaconf import DictConfig |
|
|
| from yarr.agents.agent import Agent |
| from yarr.agents.agent import ActResult |
| from yarr.agents.agent import Summary |
| from yarr.agents.agent import ScalarSummary |
|
|
|
|
| from torch.nn.parallel import DistributedDataParallel as DDP |
|
|
| from helpers.preprocess_agent import PreprocessAgent |
|
|
|
|
| from rvt.mvt.mvt import MVT |
| from rvt.models import rvt_agent |
| from rvt.utils.peract_utils import ( |
| CAMERAS, |
| SCENE_BOUNDS, |
| IMAGE_SIZE, |
| DATA_FOLDER, |
| ) |
|
|
|
|
| import rvt.config as exp_cfg_mod |
| import rvt.models.rvt_agent as rvt_agent |
| import rvt.mvt.config as mvt_cfg_mod |
|
|
|
|
| def create_agent(cfg: DictConfig): |
| exp_cfg = exp_cfg_mod.get_cfg_defaults() |
| exp_cfg.bs = cfg.replay.batch_size |
| exp_cfg.tasks = ",".join(cfg.rlbench.tasks) |
|
|
| exp_cfg.freeze() |
|
|
| mvt_cfg = mvt_cfg_mod.get_cfg_defaults() |
| mvt_cfg.proprio_dim = cfg.method.low_dim_size |
| mvt_cfg.freeze() |
|
|
| agent = RVTAgentWrapper( |
| cfg.framework.checkpoint_name_prefix, cfg.rlbench, mvt_cfg, exp_cfg |
| ) |
|
|
| preprocess_agent = PreprocessAgent(pose_agent=agent) |
| return preprocess_agent |
|
|
|
|
| class RVTAgentWrapper(Agent): |
| def __init__(self, checkpoint_name_prefix, rlbench_cfg, mvt_cfg, exp_cfg): |
| self._checkpoint_filename = f"{checkpoint_name_prefix}.pt" |
| self.rvt_agent = None |
| self.rlbench_cfg = rlbench_cfg |
| self.mvt_cfg = mvt_cfg |
| self.exp_cfg = exp_cfg |
| self._summaries = {} |
|
|
| def build(self, training: bool, device=None) -> None: |
| import torch |
|
|
| torch.cuda.set_device(device) |
| torch.cuda.empty_cache() |
|
|
| if isinstance(device, int): |
| device = f"cuda:{device}" |
|
|
| rvt = MVT( |
| renderer_device=device, |
| **self.mvt_cfg, |
| ) |
| rvt = rvt.to(device) |
|
|
| if training: |
| rvt = DDP(rvt, device_ids=[device]) |
|
|
| self.rvt_agent = rvt_agent.RVTAgent( |
| network=rvt, |
| |
| add_lang=self.mvt_cfg.add_lang, |
| scene_bounds=self.rlbench_cfg.scene_bounds, |
| cameras=self.rlbench_cfg.cameras, |
| log_dir="/tmp/eval_run", |
| **self.exp_cfg.peract, |
| **self.exp_cfg.rvt, |
| ) |
|
|
| self.rvt_agent.build(training, device) |
|
|
| def update(self, step: int, replay_sample: dict) -> dict: |
| for k, v in replay_sample.items(): |
| replay_sample[k] = v.unsqueeze(1) |
| |
| replay_sample["lang_goal_embs"] = replay_sample["lang_token_embs"] |
| replay_sample["tasks"] = self.exp_cfg.tasks.split(",") |
|
|
| update_dict = self.rvt_agent.update(step, replay_sample) |
|
|
| for key, val in self.rvt_agent.loss_log.items(): |
| self._summaries[key] = np.mean(np.array(val)) |
|
|
| return { |
| "total_losses": update_dict["total_loss"], |
| } |
|
|
| return result |
|
|
| def act(self, step: int, observation: dict, deterministic: bool) -> ActResult: |
| return self.rvt_agent.act(step, observation, deterministic) |
|
|
| def reset(self) -> None: |
| self.rvt_agent.reset() |
|
|
| def update_summaries(self) -> List[Summary]: |
| summaries = [] |
| for k, v in self._summaries.items(): |
| summaries.append(ScalarSummary(f"RVT/{k}", v)) |
| return summaries |
|
|
| def act_summaries(self) -> List[Summary]: |
| return [] |
|
|
| def load_weights(self, savedir: str) -> None: |
| """ |
| copied from RVT |
| """ |
| device = torch.device("cuda:0") |
| weight_file = os.path.join(savedir, self._checkpoint_filename) |
| state_dict = torch.load(weight_file, map_location=device) |
|
|
| model = self.rvt_agent._network |
| optimizer = self.rvt_agent._optimizer |
| lr_sched = self.rvt_agent._lr_sched |
|
|
| if isinstance(model, DDP): |
| model = model.module |
|
|
| model.load_state_dict(state_dict["model_state"]) |
| optimizer.load_state_dict(state_dict["optimizer_state"]) |
| lr_sched.load_state_dict(state_dict["lr_sched_state"]) |
|
|
| return self.rvt_agent.load_clip() |
|
|
| def save_weights(self, savedir: str) -> None: |
| os.makedirs(savedir, exist_ok=True) |
|
|
| weight_file = os.path.join(savedir, self._checkpoint_filename) |
|
|
| model = self.rvt_agent._network |
| optimizer = self.rvt_agent._optimizer |
| lr_sched = self.rvt_agent._lr_sched |
|
|
| if isinstance(model, DDP): |
| model = model.module |
|
|
| model_state = model.state_dict() |
|
|
| torch.save( |
| { |
| "model_state": model_state, |
| "optimizer_state": optimizer.state_dict(), |
| "lr_sched_state": lr_sched.state_dict(), |
| }, |
| weight_file, |
| ) |
|
|
