example/isaacgym/train_isaacgym_remote_server.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this cell only when you are troubled with the following error especially if you are using conda to manage python environment\n",
    "# `ImportError: libpython3.8.so.1.0: cannot open shared object file: No such file or directory\n",
    "import os\n",
    "from pathlib import Path\n",
    "from ctypes import cdll\n",
    "import sys\n",
    "\n",
    "is_conda = 'CONDA_PREFIX' in os.environ or 'CONDA_DEFAULT_ENV' in os.environ\n",
    "if is_conda:\n",
    "    version_info = sys.version_info\n",
    "    if version_info.major == 3 and version_info.minor >= 8:\n",
    "        conda_lib_path = Path(\n",
    "            sys.executable).parent.parent / f\"lib/libpython{version_info.major}.{version_info.minor}.so.1.0\"\n",
    "    else:\n",
    "        conda_lib_path = Path(\n",
    "            sys.executable).parent.parent / f\"lib/libpython{version_info.major}.{version_info.minor}m.so.1.0\"\n",
    "    python_lib = cdll.LoadLibrary(str(conda_lib_path))\n",
    "    print(f\"Load Python lib {conda_lib_path}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from datetime import datetime\n",
    "import numpy as np\n",
    "\n",
    "# noinspection PyUnresolvedReferences\n",
    "import isaacgym\n",
    "from sim_web_visualizer.isaac_visualizer_client import create_isaac_visualizer, bind_visualizer_to_gym, set_gpu_pipeline\n",
    "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n",
    "\n",
    "sys.path.append(os.path.join(os.path.abspath(''), 'IsaacGymEnvs'))\n",
    "\n",
    "import hydra\n",
    "\n",
    "from isaacgymenvs.pbt.pbt import PbtAlgoObserver, initial_pbt_check\n",
    "from hydra.utils import to_absolute_path\n",
    "from isaacgymenvs.tasks import isaacgym_task_map\n",
    "from omegaconf import DictConfig, OmegaConf\n",
    "import gym\n",
    "\n",
    "from isaacgymenvs.utils.reformat import omegaconf_to_dict, print_dict\n",
    "from isaacgymenvs.utils.utils import set_np_formatting, set_seed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is equivalent to run `meshcat-server` inside a separate terminal window\n",
    "start_zmq_server_as_subprocess()\n",
    "\n",
    "################################################################################################\n",
    "# The following code is only used for visualizer and not presented in the original IsaacGymEnv\n",
    "################################################################################################\n",
    "from isaacgymenvs.tasks.base import vec_task\n",
    "from isaacgym import gymapi\n",
    "\n",
    "\n",
    "def wrapped_create_sim(self: vec_task.VecTask, compute_device: int, graphics_device: int, physics_engine,\n",
    "                       sim_params: gymapi.SimParams):\n",
    "    sim = vec_task._create_sim_once(self.gym, compute_device, graphics_device, physics_engine, sim_params)\n",
    "    if sim is None:\n",
    "        print(\"*** Failed to create sim\")\n",
    "        quit()\n",
    "    self.gym = bind_visualizer_to_gym(self.gym, sim)\n",
    "    set_gpu_pipeline(sim_params.use_gpu_pipeline)\n",
    "    return sim\n",
    "\n",
    "\n",
    "# Reload VecTask function to create a hook for sim_web_visualizer\n",
    "vec_task.VecTask.create_sim = wrapped_create_sim\n",
    "\n",
    "# Create web visualizer\n",
    "web_visualizer = create_isaac_visualizer(port=6000, host=\"localhost\", keep_default_viewer=False, max_env=16,\n",
    "                                         scene_offset=np.array([40.0, 40.0]))\n",
    "\n",
    "\n",
    "################################################################################################\n",
    "# End of visualizer code\n",
    "################################################################################################\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "################################################################################################\n",
    "# The following code is the same as the original IsaacGymEnv\n",
    "# https://github.com/NVIDIA-Omniverse/IsaacGymEnvs/blob/main/isaacgymenvs/train.py\n",
    "################################################################################################\n",
    "\n",
    "def preprocess_train_config(cfg, config_dict):\n",
    "    \"\"\"\n",
    "    Adding common configuration parameters to the rl_games train config.\n",
    "    An alternative to this is inferring them in task-specific .yaml files, but that requires repeating the same\n",
    "    variable interpolations in each config.\n",
    "    \"\"\"\n",
    "\n",
    "    train_cfg = config_dict['params']['config']\n",
    "\n",
    "    train_cfg['device'] = cfg.rl_device\n",
    "\n",
    "    train_cfg['population_based_training'] = cfg.pbt.enabled\n",
    "    train_cfg['pbt_idx'] = cfg.pbt.policy_idx if cfg.pbt.enabled else None\n",
    "\n",
    "    train_cfg['full_experiment_name'] = cfg.get('full_experiment_name')\n",
    "\n",
    "    print(f'Using rl_device: {cfg.rl_device}')\n",
    "    print(f'Using sim_device: {cfg.sim_device}')\n",
    "    print(train_cfg)\n",
    "\n",
    "    try:\n",
    "        model_size_multiplier = config_dict['params']['network']['mlp']['model_size_multiplier']\n",
    "        if model_size_multiplier != 1:\n",
    "            units = config_dict['params']['network']['mlp']['units']\n",
    "            for i, u in enumerate(units):\n",
    "                units[i] = u * model_size_multiplier\n",
    "            print(\n",
    "                f'Modified MLP units by x{model_size_multiplier} to {config_dict[\"params\"][\"network\"][\"mlp\"][\"units\"]}')\n",
    "    except KeyError:\n",
    "        pass\n",
    "\n",
    "    return config_dict\n",
    "\n",
    "\n",
    "def launch_rlg_hydra(cfg: DictConfig):\n",
    "    if cfg.pbt.enabled:\n",
    "        initial_pbt_check(cfg)\n",
    "\n",
    "    from isaacgymenvs.utils.rlgames_utils import RLGPUEnv, RLGPUAlgoObserver, MultiObserver, ComplexObsRLGPUEnv\n",
    "    from isaacgymenvs.utils.wandb_utils import WandbAlgoObserver\n",
    "    from rl_games.common import env_configurations, vecenv\n",
    "    from rl_games.torch_runner import Runner\n",
    "    from rl_games.algos_torch import model_builder\n",
    "    from isaacgymenvs.learning import amp_continuous\n",
    "    from isaacgymenvs.learning import amp_players\n",
    "    from isaacgymenvs.learning import amp_models\n",
    "    from isaacgymenvs.learning import amp_network_builder\n",
    "    import isaacgymenvs\n",
    "\n",
    "    time_str = datetime.now().strftime(\"%Y-%m-%d_%H-%M-%S\")\n",
    "    run_name = f\"{cfg.wandb_name}_{time_str}\"\n",
    "\n",
    "    # ensure checkpoints can be specified as relative paths\n",
    "    if cfg.checkpoint:\n",
    "        cfg.checkpoint = to_absolute_path(cfg.checkpoint)\n",
    "\n",
    "    cfg_dict = omegaconf_to_dict(cfg)\n",
    "    print_dict(cfg_dict)\n",
    "\n",
    "    # set numpy formatting for printing only\n",
    "    set_np_formatting()\n",
    "\n",
    "    # global rank of the GPU\n",
    "    global_rank = int(os.getenv(\"RANK\", \"0\"))\n",
    "\n",
    "    # sets seed. if seed is -1 will pick a random one\n",
    "    cfg.seed = set_seed(cfg.seed, torch_deterministic=cfg.torch_deterministic, rank=global_rank)\n",
    "\n",
    "    def create_isaacgym_env(**kwargs):\n",
    "        envs = isaacgymenvs.make(\n",
    "            cfg.seed,\n",
    "            cfg.task_name,\n",
    "            cfg.task.env.numEnvs,\n",
    "            cfg.sim_device,\n",
    "            cfg.rl_device,\n",
    "            cfg.graphics_device_id,\n",
    "            cfg.headless,\n",
    "            cfg.multi_gpu,\n",
    "            cfg.capture_video,\n",
    "            cfg.force_render,\n",
    "            cfg,\n",
    "            **kwargs,\n",
    "        )\n",
    "        if cfg.capture_video:\n",
    "            envs.is_vector_env = True\n",
    "            envs = gym.wrappers.RecordVideo(\n",
    "                envs,\n",
    "                f\"videos/{run_name}\",\n",
    "                step_trigger=lambda step: step % cfg.capture_video_freq == 0,\n",
    "                video_length=cfg.capture_video_len,\n",
    "            )\n",
    "        return envs\n",
    "\n",
    "    env_configurations.register('rlgpu', {\n",
    "        'vecenv_type': 'RLGPU',\n",
    "        'env_creator': lambda **kwargs: create_isaacgym_env(**kwargs),\n",
    "    })\n",
    "\n",
    "    ige_env_cls = isaacgym_task_map[cfg.task_name]\n",
    "    dict_cls = ige_env_cls.dict_obs_cls if hasattr(ige_env_cls, 'dict_obs_cls') and ige_env_cls.dict_obs_cls else False\n",
    "\n",
    "    if dict_cls:\n",
    "\n",
    "        obs_spec = {}\n",
    "        actor_net_cfg = cfg.train.params.network\n",
    "        obs_spec['obs'] = {'names': list(actor_net_cfg.inputs.keys()),\n",
    "                           'concat': not actor_net_cfg.name == \"complex_net\", 'space_name': 'observation_space'}\n",
    "        if \"central_value_config\" in cfg.train.params.config:\n",
    "            critic_net_cfg = cfg.train.params.config.central_value_config.network\n",
    "            obs_spec['states'] = {'names': list(critic_net_cfg.inputs.keys()),\n",
    "                                  'concat': not critic_net_cfg.name == \"complex_net\", 'space_name': 'state_space'}\n",
    "\n",
    "        vecenv.register('RLGPU',\n",
    "                        lambda config_name, num_actors, **kwargs: ComplexObsRLGPUEnv(config_name, num_actors, obs_spec,\n",
    "                                                                                     **kwargs))\n",
    "    else:\n",
    "\n",
    "        vecenv.register('RLGPU', lambda config_name, num_actors, **kwargs: RLGPUEnv(config_name, num_actors, **kwargs))\n",
    "\n",
    "    rlg_config_dict = omegaconf_to_dict(cfg.train)\n",
    "    rlg_config_dict = preprocess_train_config(cfg, rlg_config_dict)\n",
    "\n",
    "    observers = [RLGPUAlgoObserver()]\n",
    "\n",
    "    if cfg.pbt.enabled:\n",
    "        pbt_observer = PbtAlgoObserver(cfg)\n",
    "        observers.append(pbt_observer)\n",
    "\n",
    "    if cfg.wandb_activate:\n",
    "        cfg.seed += global_rank\n",
    "        if global_rank == 0:\n",
    "            # initialize wandb only once per multi-gpu run\n",
    "            wandb_observer = WandbAlgoObserver(cfg)\n",
    "            observers.append(wandb_observer)\n",
    "\n",
    "    # register new AMP network builder and agent\n",
    "    def build_runner(algo_observer):\n",
    "        runner = Runner(algo_observer)\n",
    "        runner.algo_factory.register_builder('amp_continuous', lambda **kwargs: amp_continuous.AMPAgent(**kwargs))\n",
    "        runner.player_factory.register_builder('amp_continuous',\n",
    "                                               lambda **kwargs: amp_players.AMPPlayerContinuous(**kwargs))\n",
    "        model_builder.register_model('continuous_amp', lambda network, **kwargs: amp_models.ModelAMPContinuous(network))\n",
    "        model_builder.register_network('amp', lambda **kwargs: amp_network_builder.AMPBuilder())\n",
    "\n",
    "        return runner\n",
    "\n",
    "    # convert CLI arguments into dictionary\n",
    "    # create runner and set the settings\n",
    "    runner = build_runner(MultiObserver(observers))\n",
    "    runner.load(rlg_config_dict)\n",
    "    runner.reset()\n",
    "\n",
    "    # dump config dict\n",
    "    if not cfg.test:\n",
    "        experiment_dir = os.path.join('runs', cfg.train.params.config.name +\n",
    "                                      '_{date:%d-%H-%M-%S}'.format(date=datetime.now()))\n",
    "\n",
    "        os.makedirs(experiment_dir, exist_ok=True)\n",
    "        with open(os.path.join(experiment_dir, 'config.yaml'), 'w') as f:\n",
    "            f.write(OmegaConf.to_yaml(cfg))\n",
    "\n",
    "    runner.run({\n",
    "        'train': not cfg.test,\n",
    "        'play': cfg.test,\n",
    "        'checkpoint': cfg.checkpoint,\n",
    "        'sigma': cfg.sigma if cfg.sigma != '' else None\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Comment out the hydra/job_logging in the config to avoid a hydra error.\n",
    "# This is nothing related to the web visualizer but will raise a multiple value error without the following modification\n",
    "# You do not need to do that if you are not using hydra compose API.\n",
    "!sed -i 's/- hydra/#- hydra/g' IsaacGymEnvs/isaacgymenvs/cfg/config.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Modify the IP if the jupyter server is not localhost\n",
    "# If your remote server has a public ip, the machine ip should be that one\n",
    "# If you are using ssh port forwarding like: ssh YOUR_SERVER -L 7000:localhost:7000 to forward port 7000 from your server to local machine, it should still be your localhost\n",
    "machine_ip = \"127.0.0.1\"\n",
    "web_visualizer.jupyter_cell(url=f\"http://{machine_ip}:7000/static/\", height=400)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "hydra.core.global_hydra.GlobalHydra.instance().clear()\n",
    "config_path = \"IsaacGymEnvs/isaacgymenvs/cfg\"\n",
    "task_name = \"Ant\"\n",
    "with hydra.initialize(version_base=None, config_path=config_path):\n",
    "    cfg = hydra.compose(config_name='config.yaml', overrides=[f\"task={task_name}\"])\n",
    "launch_rlg_hydra(cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}