rename go1_gym and go1_gym_learn

go2_gym/envs/base/legged_robot.py

@@ -9,10 +9,10 @@ from isaacgym.torch_utils import *
 assert gymtorch
 import torch
 
-from go1_gym import MINI_GYM_ROOT_DIR
-from go1_gym.envs.base.base_task import BaseTask
-from go1_gym.utils.math_utils import quat_apply_yaw, wrap_to_pi, get_scale_shift
-from go1_gym.utils.terrain import Terrain
+from go2_gym import MINI_GYM_ROOT_DIR
+from go2_gym.envs.base.base_task import BaseTask
+from go2_gym.utils.math_utils import quat_apply_yaw, wrap_to_pi, get_scale_shift
+from go2_gym.utils.terrain import Terrain
 from .legged_robot_config import Cfg
 
 
@@ -1387,7 +1387,7 @@ class LeggedRobot(BaseTask):
             Looks for self._reward_<REWARD_NAME>, where <REWARD_NAME> are names of all non zero reward scales in the cfg.
         """
         # reward containers
-        from go1_gym.envs.rewards.corl_rewards import CoRLRewards
+        from go2_gym.envs.rewards.corl_rewards import CoRLRewards
         reward_containers = {"CoRLRewards": CoRLRewards}
         self.reward_container = reward_containers[self.cfg.rewards.reward_container_name](self)
 

go2_gym/envs/go1/go1_config.py

@@ -2,7 +2,7 @@ from typing import Union
 
 from params_proto import Meta
 
-from go1_gym.envs.base.legged_robot_config import Cfg
+from go2_gym.envs.base.legged_robot_config import Cfg
 
 
 def config_go1(Cnfg: Union[Cfg, Meta]):

go2_gym/envs/go1/velocity_tracking/__init__.py

@@ -3,8 +3,8 @@ import torch
 from params_proto import Meta
 from typing import Union
 
-from go1_gym.envs.base.legged_robot import LeggedRobot
-from go1_gym.envs.base.legged_robot_config import Cfg
+from go2_gym.envs.base.legged_robot import LeggedRobot
+from go2_gym.envs.base.legged_robot_config import Cfg
 
 
 class VelocityTrackingEasyEnv(LeggedRobot):

go2_gym/envs/rewards/corl_rewards.py

@@ -1,6 +1,6 @@
 import torch
 import numpy as np
-from go1_gym.utils.math_utils import quat_apply_yaw, wrap_to_pi, get_scale_shift
+from go2_gym.utils.math_utils import quat_apply_yaw, wrap_to_pi, get_scale_shift
 from isaacgym.torch_utils import *
 from isaacgym import gymapi
 

go2_gym/envs/wrappers/history_wrapper.py

@@ -47,12 +47,12 @@ if __name__ == "__main__":
 
     import ml_logger as logger
 
-    from go1_gym_learn.ppo import Runner
-    from go1_gym.envs.wrappers.history_wrapper import HistoryWrapper
-    from go1_gym_learn.ppo.actor_critic import AC_Args
+    from go2_gym_learn.ppo import Runner
+    from go2_gym.envs.wrappers.history_wrapper import HistoryWrapper
+    from go2_gym_learn.ppo.actor_critic import AC_Args
 
-    from go1_gym.envs.base.legged_robot_config import Cfg
-    from go1_gym.envs.mini_cheetah.mini_cheetah_config import config_mini_cheetah
+    from go2_gym.envs.base.legged_robot_config import Cfg
+    from go2_gym.envs.mini_cheetah.mini_cheetah_config import config_mini_cheetah
     config_mini_cheetah(Cfg)
 
     test_env = gym.make("VelocityTrackingEasyEnv-v0", cfg=Cfg)

go2_gym/utils/terrain.py

@@ -6,7 +6,7 @@ import numpy as np
 from isaacgym import terrain_utils
 from numpy.random import choice
 
-from go1_gym.envs.base.legged_robot_config import Cfg
+from go2_gym.envs.base.legged_robot_config import Cfg
 
 
 class Terrain:

go2_gym_learn/eval_metrics/domain_randomization.py

@@ -1,4 +1,4 @@
-from go1_gym.envs.base.legged_robot_config import Cfg
+from go2_gym.envs.base.legged_robot_config import Cfg
 
 
 def base_set():

go2_gym_learn/ppo/__init__.py

@@ -12,7 +12,7 @@ import copy
 from .actor_critic import ActorCritic
 from .rollout_storage import RolloutStorage
 
-from go1_gym import MINI_GYM_ROOT_DIR
+from go2_gym import MINI_GYM_ROOT_DIR
 
 
 def class_to_dict(obj) -> dict:
@@ -35,7 +35,7 @@ def class_to_dict(obj) -> dict:
 
 class DataCaches:
     def __init__(self, curriculum_bins):
-        from go1_gym_learn.ppo.metrics_caches import DistCache, SlotCache
+        from go2_gym_learn.ppo.metrics_caches import DistCache, SlotCache
 
         self.slot_cache = SlotCache(curriculum_bins)
         self.dist_cache = DistCache()

go2_gym_learn/ppo/ppo.py

@@ -7,9 +7,9 @@ import torch.nn.functional as F
 import numpy as np
 from params_proto import PrefixProto
 
-from go1_gym_learn.ppo import ActorCritic
-from go1_gym_learn.ppo import RolloutStorage
-from go1_gym_learn.ppo import caches
+from go2_gym_learn.ppo import ActorCritic
+from go2_gym_learn.ppo import RolloutStorage
+from go2_gym_learn.ppo import caches
 
 
 class PPO_Args(PrefixProto):

go2_gym_learn/ppo/rollout_storage.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from go1_gym_learn.utils import split_and_pad_trajectories
+from go2_gym_learn.utils import split_and_pad_trajectories
 
 class RolloutStorage:
     class Transition:

go2_gym_learn/ppo_cse/__init__.py

@@ -31,7 +31,7 @@ def class_to_dict(obj) -> dict:
 
 class DataCaches:
     def __init__(self, curriculum_bins):
-        from go1_gym_learn.ppo.metrics_caches import SlotCache, DistCache
+        from go2_gym_learn.ppo.metrics_caches import SlotCache, DistCache
 
         self.slot_cache = SlotCache(curriculum_bins)
         self.dist_cache = DistCache()

go2_gym_learn/ppo_cse/ppo.py

@@ -5,9 +5,9 @@ import torch.nn.functional as F
 import numpy as np
 from params_proto import PrefixProto
 
-from go1_gym_learn.ppo_cse import ActorCritic
-from go1_gym_learn.ppo_cse import RolloutStorage
-from go1_gym_learn.ppo_cse import caches
+from go2_gym_learn.ppo_cse import ActorCritic
+from go2_gym_learn.ppo_cse import RolloutStorage
+from go2_gym_learn.ppo_cse import caches
 
 
 class PPO_Args(PrefixProto):

go2_gym_learn/ppo_cse/rollout_storage.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from go1_gym_learn.utils import split_and_pad_trajectories
+from go2_gym_learn.utils import split_and_pad_trajectories
 
 class RolloutStorage:
     class Transition:

scripts/play.py

@@ -7,10 +7,10 @@ import numpy as np
 import glob
 import pickle as pkl
 
-from go1_gym.envs import *
-from go1_gym.envs.base.legged_robot_config import Cfg
-from go1_gym.envs.go1.go1_config import config_go1
-from go1_gym.envs.go1.velocity_tracking import VelocityTrackingEasyEnv
+from go2_gym.envs import *
+from go2_gym.envs.base.legged_robot_config import Cfg
+from go2_gym.envs.go1.go1_config import config_go1
+from go2_gym.envs.go1.velocity_tracking import VelocityTrackingEasyEnv
 
 from tqdm import tqdm
 
@@ -74,14 +74,14 @@ def load_env(label, headless=False):
     Cfg.asset.flip_visual_attachments = True
 
 
-    from go1_gym.envs.wrappers.history_wrapper import HistoryWrapper
+    from go2_gym.envs.wrappers.history_wrapper import HistoryWrapper
 
     env = VelocityTrackingEasyEnv(sim_device='cuda:0', headless=False, cfg=Cfg)
     env = HistoryWrapper(env)
 
     # load policy
     from ml_logger import logger
-    from go1_gym_learn.ppo_cse.actor_critic import ActorCritic
+    from go2_gym_learn.ppo_cse.actor_critic import ActorCritic
 
     policy = load_policy(logdir)
 
@@ -92,7 +92,7 @@ def play_go1(headless=True):
     from ml_logger import logger
 
     from pathlib import Path
-    from go1_gym import MINI_GYM_ROOT_DIR
+    from go2_gym import MINI_GYM_ROOT_DIR
     import glob
     import os
 

scripts/test.py

@@ -5,10 +5,10 @@ import matplotlib.pyplot as plt
 import torch
 from tqdm import trange
 
-from go1_gym.envs import *
-from go1_gym.envs.base.legged_robot_config import Cfg
-from go1_gym.envs.go1.go1_config import config_go1
-from go1_gym.envs.go1.velocity_tracking import VelocityTrackingEasyEnv
+from go2_gym.envs import *
+from go2_gym.envs.base.legged_robot_config import Cfg
+from go2_gym.envs.go1.go1_config import config_go1
+from go2_gym.envs.go1.velocity_tracking import VelocityTrackingEasyEnv
 
 
 def run_env(render=False, headless=False):

scripts/train.py

@@ -4,17 +4,17 @@ def train_go1(headless=True):
     assert isaacgym
     import torch
 
-    from go1_gym.envs.base.legged_robot_config import Cfg
-    from go1_gym.envs.go1.go1_config import config_go1
-    from go1_gym.envs.go1.velocity_tracking import VelocityTrackingEasyEnv
+    from go2_gym.envs.base.legged_robot_config import Cfg
+    from go2_gym.envs.go1.go1_config import config_go1
+    from go2_gym.envs.go1.velocity_tracking import VelocityTrackingEasyEnv
 
     from ml_logger import logger
 
-    from go1_gym_learn.ppo_cse import Runner
-    from go1_gym.envs.wrappers.history_wrapper import HistoryWrapper
-    from go1_gym_learn.ppo_cse.actor_critic import AC_Args
-    from go1_gym_learn.ppo_cse.ppo import PPO_Args
-    from go1_gym_learn.ppo_cse import RunnerArgs
+    from go2_gym_learn.ppo_cse import Runner
+    from go2_gym.envs.wrappers.history_wrapper import HistoryWrapper
+    from go2_gym_learn.ppo_cse.actor_critic import AC_Args
+    from go2_gym_learn.ppo_cse.ppo import PPO_Args
+    from go2_gym_learn.ppo_cse import RunnerArgs
 
     config_go1(Cfg)
 
@@ -220,7 +220,7 @@ def train_go1(headless=True):
 if __name__ == '__main__':
     from pathlib import Path
     from ml_logger import logger
-    from go1_gym import MINI_GYM_ROOT_DIR
+    from go2_gym import MINI_GYM_ROOT_DIR
 
     stem = Path(__file__).stem
     logger.configure(logger.utcnow(f'gait-conditioned-agility/%Y-%m-%d/{stem}/%H%M%S.%f'),