diff --git a/examples/go2_pos_command/README.md b/examples/go2_pos_command/README.md new file mode 100644 index 0000000..347f799 --- /dev/null +++ b/examples/go2_pos_command/README.md @@ -0,0 +1,38 @@ +# Go2 Simple Locomotion Example + +A simple program that teaches the Go2 robot to go to a target position in the world frame. + +This example uses the Genesis Forge managed environment setup, which let's the environment be dedicated more to the scene setup +and reward shaping, than logic to handle domain randomization and logging. + +## Training + +This will be trained using the [rsl_rl](https://github.com/leggedrobotics/rsl_rl) training library. So first, we need to install that and tensorboard: + +```bash +pip install tensorboard rsl-rl-lib>=2.2.4 +``` + +Now you can run the training with: + +```bash +python ./train.py +``` + + +You can view the training progress with: + +```bash +tensorboard --logdir ./logs/ +``` + +The Genesis Forge training environment will also save videos while training that can be viewed in `./logs/go2-pos-command/videos`. + + +## Evaluation + +Now you can view the trained policy: + +```bash +python ./eval.py ./logs/go2-walking/ +``` diff --git a/examples/go2_pos_command/environment.py b/examples/go2_pos_command/environment.py new file mode 100644 index 0000000..9575ed2 --- /dev/null +++ b/examples/go2_pos_command/environment.py @@ -0,0 +1,237 @@ +""" +Simplified Go2 Locomotion Environment using managers to handle everything. +""" + +import torch +import genesis as gs + +from genesis_forge import ManagedEnvironment +from genesis_forge.managers import ( + RewardManager, + TerminationManager, + EntityManager, + ObservationManager, + ActuatorManager, + PositionActionManager, + PositionCommandManager +) +from genesis_forge.mdp import reset, rewards, terminations + + +INITIAL_BODY_POSITION = [0.0, 0.0, 0.4] +INITIAL_QUAT = [1.0, 0.0, 0.0, 0.0] +TARGET_X_VELOCITY = 0.5 + + +class Go2PosEnv(ManagedEnvironment): + """ + Example training environment for the Go2 robot. + """ + + def __init__( + self, + num_envs: int = 1, + dt: float = 1 / 50, # control frequency on real robot is 50hz + max_episode_length_s: int | None = 20, + headless: bool = True, + ): + super().__init__( + num_envs=num_envs, + dt=dt, + max_episode_length_sec=max_episode_length_s, + max_episode_random_scaling=0.1, + ) + + # Construct the scene + self.scene = gs.Scene( + show_viewer=not headless, + sim_options=gs.options.SimOptions(dt=self.dt, substeps=2), + viewer_options=gs.options.ViewerOptions( + max_FPS=int(0.5 / self.dt), + camera_pos=(2.0, 0.0, 2.5), + camera_lookat=(0.0, 0.0, 0.5), + camera_fov=40, + ), + vis_options=gs.options.VisOptions(rendered_envs_idx=list(range(1))), + rigid_options=gs.options.RigidOptions( + dt=self.dt, + constraint_solver=gs.constraint_solver.Newton, + enable_collision=True, + enable_joint_limit=True, + # for this locomotion policy there are usually no more than 30 collision pairs + # set a low value can save memory + max_collision_pairs=30, + ), + ) + + # Create terrain + self.terrain = self.scene.add_entity(gs.morphs.Plane()) + + # Robot + self.robot = self.scene.add_entity( + gs.morphs.URDF( + file="urdf/go2/urdf/go2.urdf", + pos=INITIAL_BODY_POSITION, + quat=INITIAL_QUAT, + ), + ) + + # Camera, for headless video recording + self.camera = self.scene.add_camera( + pos=(-2.5, -1.5, 1.0), + lookat=(0.0, 0.0, 0.0), + res=(1280, 720), + fov=40, + env_idx=0, + debug=True, + ) + + def config(self): + """ + Configure the environment managers + """ + ## + # Robot manager + # i.e. what to do with the robot when it is reset + self.robot_manager = EntityManager( + self, + entity_attr="robot", + on_reset={ + # Reset the robot's initial position + "position": { + "fn": reset.position, + "params": { + "position": INITIAL_BODY_POSITION, + "quat": INITIAL_QUAT, + "zero_velocity": True, + }, + }, + }, + ) + self.position_command_manager = PositionCommandManager( + self, + debug_visualizer=True, + range = { + "x": (-10.0, 10.0), + "y": (-10.0, 10.0), + "z": (0.29, 0.31), + } + ) + ## + # Joint Actions + self.actuator_manager = ActuatorManager( + self, + joint_names=[ + "FL_.*_joint", + "FR_.*_joint", + "RL_.*_joint", + "RR_.*_joint", + ], + default_pos={ + ".*_hip_joint": 0.0, + "FL_thigh_joint": 0.8, + "FR_thigh_joint": 0.8, + "RL_thigh_joint": 1.0, + "RR_thigh_joint": 1.0, + ".*_calf_joint": -1.5, + }, + kp=20, + kv=0.5, + ) + self.action_manager = PositionActionManager( + self, + scale=0.25, + clip=(-100.0, 100.0), + use_default_offset=True, + actuator_manager=self.actuator_manager, + ) + + ## + # Rewards + RewardManager( + self, + logging_enabled=True, + cfg={ + "tracking_lin_vel": { + "weight": 1.0, + "fn": rewards.command_tracking_position, + "params": { + "position_cmd_manager": self.position_command_manager, + "entity_manager": self.robot_manager, + }, + }, + "lin_vel_z": { + "weight": -1.0, + "fn": rewards.lin_vel_z_l2, + "params": { + "entity_manager": self.robot_manager, + }, + }, + "action_rate": { + "weight": -0.005, + "fn": rewards.action_rate_l2, + }, + "similar_to_default": { + "weight": -0.1, + "fn": rewards.dof_similar_to_default, + "params": { + "action_manager": self.action_manager, + }, + }, + }, + ) + + ## + # Termination conditions + self.termination_manager = TerminationManager( + self, + logging_enabled=True, + term_cfg={ + # The episode ended + "timeout": { + "fn": terminations.timeout, + "time_out": True, + }, + # Terminate if the robot's pitch and yaw angles are too large + "fall_over": { + "fn": terminations.bad_orientation, + "params": { + "limit_angle": 10.0, + "entity_manager": self.robot_manager, + }, + }, + }, + ) + + ## + # Observations + ObservationManager( + self, + cfg={ + "angle_velocity": { + "fn": lambda env: self.robot_manager.get_angular_velocity(), + "scale": 0.25, + }, + "linear_velocity": { + "fn": lambda env: self.robot_manager.get_linear_velocity(), + "scale": 2.0, + }, + "projected_gravity": { + "fn": lambda env: self.robot_manager.get_projected_gravity(), + }, + "dof_position": { + "fn": lambda env: self.action_manager.get_dofs_position(), + }, + "dof_velocity": { + "fn": lambda env: self.action_manager.get_dofs_velocity(), + "scale": 0.05, + }, + "actions": { + "fn": lambda env: self.action_manager.get_actions(), + }, + }, + ) + + def build(self): + super().build() + self.camera.follow_entity(self.robot) diff --git a/examples/go2_pos_command/eval.py b/examples/go2_pos_command/eval.py new file mode 100644 index 0000000..28ad7bc --- /dev/null +++ b/examples/go2_pos_command/eval.py @@ -0,0 +1,89 @@ +import os +import glob +import torch +import pickle +import argparse +from importlib import metadata +import genesis as gs + +from genesis_forge.wrappers import RslRlWrapper +from environment import Go2PosEnv + +try: + try: + if metadata.version("rsl-rl"): + raise ImportError + except metadata.PackageNotFoundError: + if metadata.version("rsl-rl-lib").startswith("1."): + raise ImportError +except (metadata.PackageNotFoundError, ImportError) as e: + raise ImportError("Please install install 'rsl-rl-lib>=2.2.4'.") from e +from rsl_rl.runners import OnPolicyRunner + +EXPERIMENT_NAME = "go2-pos-command" + +parser = argparse.ArgumentParser(add_help=True) +parser.add_argument("-d", "--device", type=str, default="gpu") +parser.add_argument("-e", "--exp_name", type=str, default=EXPERIMENT_NAME) +args = parser.parse_args() + + +def get_latest_model(log_dir: str) -> str: + """ + Get the last model from the log directory + """ + model_checkpoints = glob.glob(os.path.join(log_dir, "model_*.pt")) + if len(model_checkpoints) == 0: + print( + f"Warning: No model files found at '{log_dir}' (you might need to train more)." + ) + exit(1) + # Sort by the file with the highest number + sorted_models = sorted( + model_checkpoints, + key=lambda x: int(os.path.basename(x).split("_")[1].split(".")[0]), + ) + return sorted_models[-1] + + +def main(): + # Processor backend (GPU or CPU) + backend = gs.gpu + if args.device == "cpu": + backend = gs.cpu + torch.set_default_device("cpu") + gs.init(logging_level="warning", backend=backend) + + # Load training configuration + log_path = f"./logs/{args.exp_name}" + [cfg] = pickle.load(open(f"{log_path}/cfgs.pkl", "rb")) + model = get_latest_model(log_path) + + # Setup environment + env = Go2PosEnv(num_envs=1, headless=False) + env = RslRlWrapper(env) + env.build() + + # Eval + print("🎬 Loading last model...") + runner = OnPolicyRunner(env, cfg, log_path, device=gs.device) + runner.load(model) + policy = runner.get_inference_policy(device=gs.device) + + try: + obs, _ = env.reset() + with torch.no_grad(): + while True: + actions = policy(obs) + obs, _rews, _dones, _infos = env.step(actions) + except KeyboardInterrupt: + pass + except gs.GenesisException as e: + if e.message != "Viewer closed.": + raise e + except Exception as e: + raise e + + +if __name__ == "__main__": + main() diff --git a/examples/go2_pos_command/train.py b/examples/go2_pos_command/train.py new file mode 100644 index 0000000..48064b5 --- /dev/null +++ b/examples/go2_pos_command/train.py @@ -0,0 +1,134 @@ +import os +import copy +import torch +import shutil +import pickle +import argparse +from importlib import metadata +import genesis as gs + +from genesis_forge.wrappers import ( + VideoWrapper, + RslRlWrapper, +) +from environment import Go2PosEnv + +try: + try: + if metadata.version("rsl-rl"): + raise ImportError + except metadata.PackageNotFoundError: + if metadata.version("rsl-rl-lib").startswith("1."): + raise ImportError +except (metadata.PackageNotFoundError, ImportError) as e: + raise ImportError("Please install install 'rsl-rl-lib>=2.2.4'.") from e +from rsl_rl.runners import OnPolicyRunner + +EXPERIMENT_NAME = "go2-pos-command" + +parser = argparse.ArgumentParser(add_help=True) +parser.add_argument("-n", "--num_envs", type=int, default=4096) +parser.add_argument("--max_iterations", type=int, default=101) +parser.add_argument("-d", "--device", type=str, default="gpu") +parser.add_argument("-e", "--exp_name", type=str, default=EXPERIMENT_NAME) +args = parser.parse_args() + + +def training_cfg(exp_name: str, max_iterations: int): + return { + "algorithm": { + "class_name": "PPO", + "clip_param": 0.2, + "desired_kl": 0.01, + "entropy_coef": 0.01, + "gamma": 0.99, + "lam": 0.95, + "learning_rate": 0.001, + "max_grad_norm": 1.0, + "num_learning_epochs": 5, + "num_mini_batches": 4, + "schedule": "adaptive", + "use_clipped_value_loss": True, + "value_loss_coef": 1.0, + }, + "init_member_classes": {}, + "policy": { + "activation": "elu", + "actor_hidden_dims": [512, 256, 128], + "critic_hidden_dims": [512, 256, 128], + "init_noise_std": 1.0, + "class_name": "ActorCritic", + }, + "runner": { + "checkpoint": -1, + "experiment_name": exp_name, + "load_run": -1, + "log_interval": 1, + "max_iterations": max_iterations, + "record_interval": -1, + "resume": False, + "resume_path": None, + "run_name": "", + }, + "runner_class_name": "OnPolicyRunner", + "seed": 1, + "num_steps_per_env": 24, + "save_interval": 100, + "empirical_normalization": None, + "obs_groups": {"policy": ["policy"], "critic": ["policy"]}, + } + + +def main(): + # Initialize Genesis + # Processor backend (GPU or CPU) + backend = gs.gpu + if args.device == "cpu": + backend = gs.cpu + torch.set_default_device("cpu") + gs.init(logging_level="warning", backend=backend) + + # Logging directory + log_base_dir = "./logs" + experiment_name = args.exp_name + log_path = os.path.join(log_base_dir, experiment_name) + if os.path.exists(log_path): + shutil.rmtree(log_path) + os.makedirs(log_path, exist_ok=True) + print(f"Logging to: {log_path}") + + # Load training configuration and save snapshot of training configs + cfg = training_cfg(experiment_name, args.max_iterations) + pickle.dump( + [cfg], + open(os.path.join(log_path, "cfgs.pkl"), "wb"), + ) + + # Create environment + env = Go2PosEnv(num_envs=args.num_envs, headless=True) + + # Record videos in regular intervals + env = VideoWrapper( + env, + video_length_sec=12, + out_dir=os.path.join(log_path, "videos"), + episode_trigger=lambda episode_id: episode_id % 5 == 0, + ) + + # Build the environment + env = RslRlWrapper(env) + env.build() + env.reset() + + # Train + print("💪 Training model...") + runner = OnPolicyRunner(env, copy.deepcopy(cfg), log_path, device=gs.device) + runner.git_status_repos = ["."] + runner.learn( + num_learning_iterations=args.max_iterations, init_at_random_ep_len=False + ) + env.close() + + +if __name__ == "__main__": + main() diff --git a/examples/position_command/README.md b/examples/position_command/README.md new file mode 100644 index 0000000..347f799 --- /dev/null +++ b/examples/position_command/README.md @@ -0,0 +1,38 @@ +# Go2 Simple Locomotion Example + +A simple program that teaches the Go2 robot to go to a target position in the world frame. + +This example uses the Genesis Forge managed environment setup, which let's the environment be dedicated more to the scene setup +and reward shaping, than logic to handle domain randomization and logging. + +## Training + +This will be trained using the [rsl_rl](https://github.com/leggedrobotics/rsl_rl) training library. So first, we need to install that and tensorboard: + +```bash +pip install tensorboard rsl-rl-lib>=2.2.4 +``` + +Now you can run the training with: + +```bash +python ./train.py +``` + + +You can view the training progress with: + +```bash +tensorboard --logdir ./logs/ +``` + +The Genesis Forge training environment will also save videos while training that can be viewed in `./logs/go2-pos-command/videos`. + + +## Evaluation + +Now you can view the trained policy: + +```bash +python ./eval.py ./logs/go2-walking/ +``` diff --git a/examples/position_command/environment.py b/examples/position_command/environment.py new file mode 100644 index 0000000..3e8af7b --- /dev/null +++ b/examples/position_command/environment.py @@ -0,0 +1,244 @@ +""" +Simplified Go2 Locomotion Environment using managers to handle everything. +""" + +import torch +import genesis as gs + +from genesis_forge import ManagedEnvironment +from genesis_forge.managers import ( + RewardManager, + TerminationManager, + EntityManager, + ObservationManager, + ActuatorManager, + PositionActionManager, + PositionCommandManager +) +from genesis_forge.mdp import reset, rewards, terminations + + +INITIAL_BODY_POSITION = [0.0, 0.0, 0.4] +INITIAL_QUAT = [1.0, 0.0, 0.0, 0.0] +TARGET_X_VELOCITY = 0.5 + + +class Go2PosEnv(ManagedEnvironment): + """ + Example training environment for the Go2 robot. + """ + + def __init__( + self, + num_envs: int = 1, + dt: float = 1 / 50, # control frequency on real robot is 50hz + max_episode_length_s: int | None = 200, + headless: bool = True, + ): + super().__init__( + num_envs=num_envs, + dt=dt, + max_episode_length_sec=max_episode_length_s, + max_episode_random_scaling=0.1, + ) + + # Construct the scene + self.scene = gs.Scene( + show_viewer=not headless, + sim_options=gs.options.SimOptions(dt=self.dt, substeps=2), + viewer_options=gs.options.ViewerOptions( + max_FPS=int(0.5 / self.dt), + camera_pos=(2.0, 0.0, 2.5), + camera_lookat=(0.0, 0.0, 0.5), + camera_fov=40, + ), + vis_options=gs.options.VisOptions(rendered_envs_idx=list(range(1))), + rigid_options=gs.options.RigidOptions( + dt=self.dt, + constraint_solver=gs.constraint_solver.Newton, + enable_collision=True, + enable_joint_limit=True, + # for this locomotion policy there are usually no more than 30 collision pairs + # set a low value can save memory + max_collision_pairs=30, + ), + ) + + # Create terrain + self.terrain = self.scene.add_entity(gs.morphs.Plane()) + + # Robot + self.robot = self.scene.add_entity( + gs.morphs.URDF( + file="urdf/go2/urdf/go2.urdf", + pos=INITIAL_BODY_POSITION, + quat=INITIAL_QUAT, + ), + ) + + # Camera, for headless video recording + self.camera = self.scene.add_camera( + pos=(-2.5, -1.5, 1.0), + lookat=(0.0, 0.0, 0.0), + res=(1280, 720), + fov=40, + env_idx=0, + debug=True, + ) + + def config(self): + """ + Configure the environment managers + """ + ## + # Robot manager + # i.e. what to do with the robot when it is reset + self.robot_manager = EntityManager( + self, + entity_attr="robot", + on_reset={ + # Reset the robot's initial position + "position": { + "fn": reset.position, + "params": { + "position": INITIAL_BODY_POSITION, + "quat": INITIAL_QUAT, + "zero_velocity": True, + }, + }, + }, + ) + self.position_command_manager = PositionCommandManager( + self, + debug_visualizer=True, + range={ + "x": (-3.0, 3.0), + "y": (-3.0, 3.0), + "z": (0.29, 0.31), + }, + resample_time_sec=5.0, + ) + ## + # Joint Actions + self.actuator_manager = ActuatorManager( + self, + joint_names=[ + "FL_.*_joint", + "FR_.*_joint", + "RL_.*_joint", + "RR_.*_joint", + ], + default_pos={ + ".*_hip_joint": 0.0, + "FL_thigh_joint": 0.8, + "FR_thigh_joint": 0.8, + "RL_thigh_joint": 1.0, + "RR_thigh_joint": 1.0, + ".*_calf_joint": -1.5, + }, + kp=20, + kv=0.5, + ) + self.action_manager = PositionActionManager( + self, + scale=0.25, + clip=(-100.0, 100.0), + use_default_offset=True, + actuator_manager=self.actuator_manager, + ) + + ## + # Rewards + RewardManager( + self, + logging_enabled=True, + cfg={ + "tracking_pos": { + "weight": 5.0, + "fn": rewards.command_tracking_position, + "params": { + "position_cmd_manager": self.position_command_manager, + "entity_manager": self.robot_manager, + "sensitivity": 0.25, + }, + }, + "lin_vel_z": { + "weight": -1.0, + "fn": rewards.lin_vel_z_l2, + "params": { + "entity_manager": self.robot_manager, + }, + }, + "action_rate": { + "weight": -0.005, + "fn": rewards.action_rate_l2, + }, + "similar_to_default": { + "weight": -0.1, + "fn": rewards.dof_similar_to_default, + "params": { + "action_manager": self.action_manager, + }, + }, + }, + ) + + ## + # Termination conditions + self.termination_manager = TerminationManager( + self, + logging_enabled=True, + term_cfg={ + # The episode ended + "timeout": { + "fn": terminations.timeout, + "time_out": True, + }, + # Terminate if the robot's pitch and yaw angles are too large + "fall_over": { + "fn": terminations.bad_orientation, + "params": { + "limit_angle": 25.0, + "entity_manager": self.robot_manager, + }, + }, + }, + ) + + ## + # Observations + ObservationManager( + self, + cfg={ + "command": { + "fn": self.position_command_manager.observation, + "scale": 1.0, + }, + "position": { + "fn": lambda env: self.robot_manager.base_pos, + "scale": 1.0, + }, + "linear_velocity": { + "fn": lambda env: self.robot_manager.get_linear_velocity(), + "scale": 0.5, + }, + "projected_gravity": { + "fn": lambda env: self.robot_manager.get_projected_gravity(), + }, + "dof_position": { + "fn": lambda env: self.action_manager.get_dofs_position(), + }, + "dof_velocity": { + "fn": lambda env: self.action_manager.get_dofs_velocity(), + "scale": 0.05, + }, + "actions": { + "fn": lambda env: self.action_manager.get_actions(), + }, + }, + ) + + def build(self): + super().build() + self.camera.follow_entity(self.robot) + diff --git a/examples/position_command/eval.py b/examples/position_command/eval.py new file mode 100644 index 0000000..28ad7bc --- /dev/null +++ b/examples/position_command/eval.py @@ -0,0 +1,89 @@ +import os +import glob +import torch +import pickle +import argparse +from importlib import metadata +import genesis as gs + +from genesis_forge.wrappers import RslRlWrapper +from environment import Go2PosEnv + +try: + try: + if metadata.version("rsl-rl"): + raise ImportError + except metadata.PackageNotFoundError: + if metadata.version("rsl-rl-lib").startswith("1."): + raise ImportError +except (metadata.PackageNotFoundError, ImportError) as e: + raise ImportError("Please install install 'rsl-rl-lib>=2.2.4'.") from e +from rsl_rl.runners import OnPolicyRunner + +EXPERIMENT_NAME = "go2-pos-command" + +parser = argparse.ArgumentParser(add_help=True) +parser.add_argument("-d", "--device", type=str, default="gpu") +parser.add_argument("-e", "--exp_name", type=str, default=EXPERIMENT_NAME) +args = parser.parse_args() + + +def get_latest_model(log_dir: str) -> str: + """ + Get the last model from the log directory + """ + model_checkpoints = glob.glob(os.path.join(log_dir, "model_*.pt")) + if len(model_checkpoints) == 0: + print( + f"Warning: No model files found at '{log_dir}' (you might need to train more)." + ) + exit(1) + # Sort by the file with the highest number + sorted_models = sorted( + model_checkpoints, + key=lambda x: int(os.path.basename(x).split("_")[1].split(".")[0]), + ) + return sorted_models[-1] + + +def main(): + # Processor backend (GPU or CPU) + backend = gs.gpu + if args.device == "cpu": + backend = gs.cpu + torch.set_default_device("cpu") + gs.init(logging_level="warning", backend=backend) + + # Load training configuration + log_path = f"./logs/{args.exp_name}" + [cfg] = pickle.load(open(f"{log_path}/cfgs.pkl", "rb")) + model = get_latest_model(log_path) + + # Setup environment + env = Go2PosEnv(num_envs=1, headless=False) + env = RslRlWrapper(env) + env.build() + + # Eval + print("🎬 Loading last model...") + runner = OnPolicyRunner(env, cfg, log_path, device=gs.device) + runner.load(model) + policy = runner.get_inference_policy(device=gs.device) + + try: + obs, _ = env.reset() + with torch.no_grad(): + while True: + actions = policy(obs) + obs, _rews, _dones, _infos = env.step(actions) + except KeyboardInterrupt: + pass + except gs.GenesisException as e: + if e.message != "Viewer closed.": + raise e + except Exception as e: + raise e + + +if __name__ == "__main__": + main() diff --git a/examples/position_command/train.py b/examples/position_command/train.py new file mode 100644 index 0000000..6590d1a --- /dev/null +++ b/examples/position_command/train.py @@ -0,0 +1,134 @@ +import os +import copy +import torch +import shutil +import pickle +import argparse +from importlib import metadata +import genesis as gs + +from genesis_forge.wrappers import ( + VideoWrapper, + RslRlWrapper, +) +from environment import Go2PosEnv + +try: + try: + if metadata.version("rsl-rl"): + raise ImportError + except metadata.PackageNotFoundError: + if metadata.version("rsl-rl-lib").startswith("1."): + raise ImportError +except (metadata.PackageNotFoundError, ImportError) as e: + raise ImportError("Please install install 'rsl-rl-lib>=2.2.4'.") from e +from rsl_rl.runners import OnPolicyRunner + +EXPERIMENT_NAME = "go2-pos-command" + +parser = argparse.ArgumentParser(add_help=True) +parser.add_argument("-n", "--num_envs", type=int, default=4096) +parser.add_argument("--max_iterations", type=int, default=400) +parser.add_argument("-d", "--device", type=str, default="gpu") +parser.add_argument("-e", "--exp_name", type=str, default=EXPERIMENT_NAME) +args = parser.parse_args() + + +def training_cfg(exp_name: str, max_iterations: int): + return { + "algorithm": { + "class_name": "PPO", + "clip_param": 0.2, + "desired_kl": 0.01, + "entropy_coef": 0.01, + "gamma": 0.99, + "lam": 0.95, + "learning_rate": 0.001, + "max_grad_norm": 1.0, + "num_learning_epochs": 5, + "num_mini_batches": 4, + "schedule": "adaptive", + "use_clipped_value_loss": True, + "value_loss_coef": 1.0, + }, + "init_member_classes": {}, + "policy": { + "activation": "elu", + "actor_hidden_dims": [512, 256, 128], + "critic_hidden_dims": [512, 256, 128], + "init_noise_std": 1.0, + "class_name": "ActorCritic", + }, + "runner": { + "checkpoint": -1, + "experiment_name": exp_name, + "load_run": -1, + "log_interval": 1, + "max_iterations": max_iterations, + "record_interval": -1, + "resume": False, + "resume_path": None, + "run_name": "", + }, + "runner_class_name": "OnPolicyRunner", + "seed": 1, + "num_steps_per_env": 24, + "save_interval": 100, + "empirical_normalization": None, + "obs_groups": {"policy": ["policy"], "critic": ["policy"]}, + } + + +def main(): + # Initialize Genesis + # Processor backend (GPU or CPU) + backend = gs.gpu + if args.device == "cpu": + backend = gs.cpu + torch.set_default_device("cpu") + gs.init(logging_level="warning", backend=backend) + + # Logging directory + log_base_dir = "./logs" + experiment_name = args.exp_name + log_path = os.path.join(log_base_dir, experiment_name) + if os.path.exists(log_path): + shutil.rmtree(log_path) + os.makedirs(log_path, exist_ok=True) + print(f"Logging to: {log_path}") + + # Load training configuration and save snapshot of training configs + cfg = training_cfg(experiment_name, args.max_iterations) + pickle.dump( + [cfg], + open(os.path.join(log_path, "cfgs.pkl"), "wb"), + ) + + # Create environment + env = Go2PosEnv(num_envs=args.num_envs, headless=True) + + # Record videos in regular intervals + env = VideoWrapper( + env, + video_length_sec=12, + out_dir=os.path.join(log_path, "videos"), + episode_trigger=lambda episode_id: episode_id % 5 == 0, + ) + + # Build the environment + env = RslRlWrapper(env) + env.build() + env.reset() + + # Train + print("💪 Training model...") + runner = OnPolicyRunner(env, copy.deepcopy(cfg), log_path, device=gs.device) + runner.git_status_repos = ["."] + runner.learn( + num_learning_iterations=args.max_iterations, init_at_random_ep_len=False + ) + env.close() + + +if __name__ == "__main__": + main() diff --git a/genesis_forge/managers/__init__.py b/genesis_forge/managers/__init__.py index db47035..16040ca 100644 --- a/genesis_forge/managers/__init__.py +++ b/genesis_forge/managers/__init__.py @@ -3,7 +3,12 @@ from .termination_manager import TerminationManager from .action.position_action_manager import PositionActionManager from .action.position_within_limits import PositionWithinLimitsActionManager -from .command import CommandManager, VelocityCommandManager +from .command import ( + CommandManager, + PositionCommandManager, + PoseCommandManager, + VelocityCommandManager, +) from .contact import ContactManager from .terrain_manager import TerrainManager from .entity_manager import EntityManager @@ -19,6 +24,8 @@ "RewardManager", "TerminationManager", "CommandManager", + "PositionCommandManager", + "PoseCommandManager", "VelocityCommandManager", "PositionActionManager", "PositionWithinLimitsActionManager", diff --git a/genesis_forge/managers/command/__init__.py b/genesis_forge/managers/command/__init__.py index 6a8040c..31f3ad0 100644 --- a/genesis_forge/managers/command/__init__.py +++ b/genesis_forge/managers/command/__init__.py @@ -1,7 +1,11 @@ from .command_manager import CommandManager +from .position_command import PositionCommandManager +from .pose_command import PoseCommandManager from .velocity_command import VelocityCommandManager __all__ = [ "CommandManager", + "PositionCommandManager", + "PoseCommandManager", "VelocityCommandManager", ] diff --git a/genesis_forge/managers/command/pose_command.py b/genesis_forge/managers/command/pose_command.py new file mode 100644 index 0000000..044453a --- /dev/null +++ b/genesis_forge/managers/command/pose_command.py @@ -0,0 +1,232 @@ +from typing import Tuple, TypedDict + +import os +import torch +import genesis as gs +from genesis.utils.geom import euler_to_R +import numpy as np +from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.gamepads import Gamepad + +from .command_manager import CommandManager, CommandRangeValue + + +class PoseCommandRange(TypedDict): + pos_x: CommandRangeValue + pos_y: CommandRangeValue + pos_z: CommandRangeValue + euler_x: CommandRangeValue + euler_y: CommandRangeValue + euler_z: CommandRangeValue + + +class PoseDebugVisualizerConfig(TypedDict): + """Defines the configuration for the debug visualizer.""" + + envs_idx: list[int] + """The indices of the environments to visualize. If None, all environments will be visualized.""" + + arrow_offset: float + """The vertical offset of the debug arrows from the top of the robot""" + + arrow_radius: float + """The radius of the shaft of the debug arrows""" + + commanded_color: Tuple[float, float, float, float] + """The color of the pose arrow""" + + +DEFAULT_VISUALIZER_CONFIG: PoseDebugVisualizerConfig = { + "envs_idx": None, + "arrow_offset": 0.03, + "arrow_radius": 0.02, + "commanded_color": (0.0, 0.5, 0.0, 1.0), +} + + +class PoseCommandManager(CommandManager): + """ + Generates a pose command from uniform distribution. + The command comprises of a (x,y,z) position and (x,y,z) euler angles + + IMPORTANT: The pose commands are interpreted as world-relative coordinates: + - pos-X: x coordinate of the target position + - pos-Y: y coordinate of the target position + - pos-Z: z coordinate of the target position + - euler-X: x coordinate of the target orientation + - euler-Y: y coordinate of the target orientation + - euler-Z: z coordinate of the target orientation + + :::{admonition} Debug Visualization + + If you set `debug_visualizer` to True, target arrow will be rendered above the target pose + + Arrow meanings: + + - GREEN: Commanded pose for the robot in the world frame + + Args: + env: The environment to control + range: The ranges of positions and orientation + resample_time_sec: The time interval between changing the command + debug_visualizer: Enable the debug arrow visualization + debug_visualizer_cfg: The configuration for the debug visualizer + + Example:: + + class MyEnv(GenesisEnv): + def config(self): + # Create a pose command manager + self.pose_command_manager = PoseCommandManager( + self, + debug_visualizer=True, + range = { + "pos_x": (-5.0, 5.0), + "pos_y": (-5.0, 5.0), + "euler_z": (-1.57, 1.57), + } + ) + + RewardManager( + self, + logging_enabled=True, + cfg={ + "tracking_pose": { + "weight": 1.0, + "fn": rewards.command_tracking_pose, + "params": { + "pose_cmd_manager": self.pose_command_manager, + }, + }, + # ... other rewards ... + }, + ) + + # Observations + ObservationManager( + self, + cfg={ + "pose_cmd": {"fn": self.pose_command_manager.observation}, + # ... other observations ... + }, + ) + """ + + def __init__( + self, + env: GenesisEnv, + range: PoseCommandRange, + resample_time_sec: float = 5.0, + debug_visualizer: bool = False, + debug_visualizer_cfg: PoseDebugVisualizerConfig = DEFAULT_VISUALIZER_CONFIG, + ): + super().__init__(env, range=range, resample_time_sec=resample_time_sec) + self._arrow_nodes: list = [] + self.debug_visualizer = debug_visualizer + self.visualizer_cfg = {**DEFAULT_VISUALIZER_CONFIG, **debug_visualizer_cfg} + self.debug_envs_idx = None + + """ + Lifecycle Operations + """ + + def build(self): + """Build the pose command manager""" + super().build() + + # If debug envs_idx is not set, attempt to use the vis_options rendered_envs_idx + if ( + not self.debug_visualizer + or self.visualizer_cfg is None + or self.env.scene is None + ): + return + self.debug_envs_idx = self.visualizer_cfg.get("envs_idx", None) + if self.debug_envs_idx is None and self.env.scene.vis_options is not None: + self.debug_envs_idx = self.env.scene.vis_options.rendered_envs_idx + if self.debug_envs_idx is None: + self.debug_envs_idx = list[int](range(self.env.num_envs)) + + def step(self): + """Render the command arrows""" + if not self.enabled: + return + super().step() + self._render_arrow() + + def use_gamepad( + self, + gamepad: Gamepad, + pos_x: int = 0, + pos_y: int = 1, + pos_z: int = 2, + euler_x: int = 3, + euler_y: int = 4, + euler_z: int = 5, + ): + """ + Use a connected gamepad to control the command. + + Args: + gamepad: The gamepad to use. + pos_x: Map this gamepad axis index to the position in the x-direction. + pos_y: Map this gamepad axis index to the position in the y-direction. + pos_z: Map this gamepad axis index to the position in the z-direction. + euler_x: Map this gamepad axis index to the orientation in the x-direction. + euler_y: Map this gamepad axis index to the orientation in the y-direction. + euler_z: Map this gamepad axis index to the orientation in the z-direction. + """ + super().use_gamepad( + gamepad, + range_axis={ + "pos_x": pos_x, + "pos_y": pos_y, + "pos_z": pos_z, + "euler_x": euler_x, + "euler_y": euler_y, + "euler_z": euler_z, + }, + ) + + """ + Internal Implementation + """ + + def _render_arrow(self): + """ + Render the command arrow showing pose commands. + + The commanded pose arrow (green) shows the pose in the world frame + """ + if not self.debug_visualizer: + return + + # Remove existing arrows + for arrow in self._arrow_nodes: + self.env.scene.clear_debug_object(arrow) + self._arrow_nodes = [] + + for i in self.debug_envs_idx: + # Target arrow (robot-relative command transformed to world coordinates for visualization) + self._draw_arrow( + pos=self.command[i], + color=self.visualizer_cfg["commanded_color"], + ) + + def _draw_arrow( + self, + pos: torch.Tensor, + euler: torch.Tensor, + color: list[float], + ): + try: + node = self.env.scene.draw_debug_arrow( + pos=pos.cpu().numpy(), + vec=np.tile([0, 0, 1], (pos.shape[0], 1)) @ euler_to_R(euler), + color=color, + radius=self.visualizer_cfg["arrow_radius"], + ) + if node: + self._sphere_nodes.append(node) + except Exception as e: + print(f"Error adding debug visualizing in PoseCommandManager: {e}") diff --git a/genesis_forge/managers/command/position_command.py b/genesis_forge/managers/command/position_command.py new file mode 100644 index 0000000..fd9242e --- /dev/null +++ b/genesis_forge/managers/command/position_command.py @@ -0,0 +1,213 @@ +from typing import Tuple, TypedDict + +import os +import torch +import genesis as gs + +from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.gamepads import Gamepad + +from .command_manager import CommandManager, CommandRangeValue + + +class PositionCommandRange(TypedDict): + x: CommandRangeValue + y: CommandRangeValue + z: CommandRangeValue + + +class PositionDebugVisualizerConfig(TypedDict): + """Defines the configuration for the debug visualizer.""" + + envs_idx: list[int] + """The indices of the environments to visualize. If None, all environments will be visualized.""" + + sphere_offset: float + """The vertical offset of the debug sphere from the top of the robot""" + + sphere_radius: float + """The radius of the shaft of the debug sphere""" + + commanded_color: Tuple[float, float, float, float] + """The color of the commanded position sphere""" + + +DEFAULT_VISUALIZER_CONFIG: PositionDebugVisualizerConfig = { + "envs_idx": None, + "sphere_offset": 0.1, + "sphere_radius": 0.02, + "commanded_color": (0.0, 0.5, 0.0, 1.0), +} + + +class PositionCommandManager(CommandManager): + """ + Generates a position command from uniform distribution. + The command comprises of a position in the x, y, and z axes. + + IMPORTANT: The position commands are interpreted as world-relative coordinates: + - X-axis: x coordinate of the target position + - Y-axis: y coordinate of the target position + - Z-axis: z coordinate of the target position + + :::{admonition} Debug Visualization + + If you set `debug_visualizer` to True, target sphere will be rendered above the target pos + + Sphere meanings: + + - GREEN: Commanded position for the robot in the world frame + + Args: + env: The environment to control + range: The ranges of linear & angular velocities + resample_time_sec: The time interval between changing the command + debug_visualizer: Enable the debug sphere visualization + debug_visualizer_cfg: The configuration for the debug visualizer + + Example:: + + class MyEnv(GenesisEnv): + def config(self): + # Create a position command manager + self.position_command_manager = PositionCommandManager( + self, + debug_visualizer=True, + range = { + "x": (-5.0, 5.0), + "y": (-5.0, 5.0), + "z": (0.29, 0.31), + } + ) + + RewardManager( + self, + logging_enabled=True, + cfg={ + "tracking_position": { + "weight": 1.0, + "fn": rewards.command_tracking_pos, + "params": { + "position_cmd_manager": self.position_command_manager, + }, + }, + # ... other rewards ... + }, + ) + + # Observations + ObservationManager( + self, + cfg={ + "position_cmd": {"fn": self.position_command_manager.observation}, + # ... other observations ... + }, + ) + """ + + def __init__( + self, + env: GenesisEnv, + range: PositionCommandRange, + resample_time_sec: float = 5.0, + debug_visualizer: bool = False, + debug_visualizer_cfg: PositionDebugVisualizerConfig = DEFAULT_VISUALIZER_CONFIG, + ): + super().__init__(env, range=range, resample_time_sec=resample_time_sec) + self._sphere_nodes = None + self.debug_visualizer = debug_visualizer + self.visualizer_cfg = {**DEFAULT_VISUALIZER_CONFIG, **debug_visualizer_cfg} + self.debug_envs_idx = None + + """ + Lifecycle Operations + """ + + def build(self): + """Build the position command manager""" + super().build() + + # If debug envs_idx is not set, attempt to use the vis_options rendered_envs_idx + if ( + not self.debug_visualizer + or self.visualizer_cfg is None + or self.env.scene is None + ): + return + self.debug_envs_idx = self.visualizer_cfg.get("envs_idx", None) + if self.debug_envs_idx is None and self.env.scene.vis_options is not None: + self.debug_envs_idx = self.env.scene.vis_options.rendered_envs_idx + if self.debug_envs_idx is None: + self.debug_envs_idx = list[int](range(self.env.num_envs)) + + def step(self): + """Render the command spheres""" + if not self.enabled: + return + super().step() + self._render_sphere() + + def use_gamepad( + self, + gamepad: Gamepad, + x: int = 0, + y: int = 1, + z: int = 2, + ): + """ + Use a connected gamepad to control the command. + + Args: + gamepad: The gamepad to use. + x: Map this gamepad axis index to the position in the x-direction. + y: Map this gamepad axis index to the position in the y-direction. + z: Map this gamepad axis index to the position in the z-direction. + """ + super().use_gamepad( + gamepad, + range_axis={ + "x": x, + "y": y, + "z": z, + }, + ) + + """ + Internal Implementation + """ + + def _render_sphere(self): + """ + Render the command sphere showing position commands. + + The commanded position sphere (green) shows the position in the world frame + """ + if not self.debug_visualizer: + return + + # Remove existing spheres + if self._sphere_nodes is not None: + self.env.scene.clear_debug_object(self._sphere_nodes) + self._sphere_nodes = None + + # Target sphere in the world frame for visualization) + self._draw_spheres( + pos=self.command[self.debug_envs_idx], + color=self.visualizer_cfg["commanded_color"], + ) + + def _draw_spheres( + self, + pos: torch.Tensor, + color: list[float], + ): + try: + nodes = self.env.scene.draw_debug_spheres( + poss=pos.cpu().numpy(), + color=color, + radius=self.visualizer_cfg["sphere_radius"], + ) + if nodes: + self._sphere_nodes=nodes + except Exception as e: + print(f"Error adding debug visualizing in PositionCommandManager: {e}") diff --git a/genesis_forge/managers/command/velocity_command.py b/genesis_forge/managers/command/velocity_command.py index 942183c..67be6da 100644 --- a/genesis_forge/managers/command/velocity_command.py +++ b/genesis_forge/managers/command/velocity_command.py @@ -175,7 +175,11 @@ def build(self): def build_debug(self): """Build the debug components of the velocity command manager""" - if not self.debug_visualizer or self.visualizer_cfg is None: + if ( + not self.debug_visualizer + or self.visualizer_cfg is None + or self.env.scene is None + ): return # Pre-allocate buffers diff --git a/genesis_forge/mdp/rewards.py b/genesis_forge/mdp/rewards.py index eae526b..83e3307 100644 --- a/genesis_forge/mdp/rewards.py +++ b/genesis_forge/mdp/rewards.py @@ -7,10 +7,13 @@ import torch import genesis as gs +from genesis.utils.geom import quat_to_xyz from genesis_forge.genesis_env import GenesisEnv from genesis_forge.managers import ( ActuatorManager, CommandManager, + PositionCommandManager, + PoseCommandManager, VelocityCommandManager, PositionActionManager, ContactManager, @@ -282,6 +285,138 @@ def action_rate_l2(env: GenesisEnv) -> torch.Tensor: return torch.sum(torch.square(last_actions - actions), dim=1) +""" +Position Command Rewards +""" + + +def command_tracking_position( + env: GenesisEnv, + command: torch.Tensor = None, + position_cmd_manager: PositionCommandManager = None, + sensitivity: float = 0.25, + entity_attr: str = "robot", + link_name: str = None, + entity_manager: EntityManager = None, +) -> torch.Tensor: + """ + Penalize base pos away from target. + + Args: + env: The Genesis environment containing the robot + command: The commanded XYZ position the the world frame, its shape is(num_envs, 3) + position_cmd_manager: The Position command manager + sensitivity: A lower value means the reward is more sensitive to the error + entity_attr: The attribute name of the entity in the environment. + entity_manager: The entity manager for the entity. + + Returns: + torch.Tensor: Penalty for base position away from target + """ + assert ( + command is not None or position_cmd_manager is not None + ), "Either command or position_cmd_manager must be provided to commonad_tracking_base_position" + + if entity_manager is not None: + entity = entity_manager.entity + else: + entity = getattr(env, entity_attr) + + if link_name is None or link_name == "": + link = entity.base_link + else: + try: + link = entity.get_link(link_name) + except Exception as e: + raise ValueError( + f"Link name '{link_name}' not found in entity '{entity_attr}'." + ) + + pos = link.get_pos() + if position_cmd_manager is not None: + command = position_cmd_manager.command + + pos_error = torch.sum(torch.square(command - pos), dim=1) + return torch.exp(-pos_error / sensitivity) + + +""" +Pose Command Rewards +""" + + +def command_tracking_pose( + env: GenesisEnv, + command: torch.Tensor = None, + pose_cmd_manager: PoseCommandManager = None, + pos_sensitivity: float = 0.25, + euler_sensitivity: float = 0.25, + entity_attr: str = "robot", + link_name: str = None, + entity_manager: "EntityManager" = None, +) -> torch.Tensor: + """ + Penalize base pose away from target (both position and orientation) with separate sensitivities. + + Args: + env: The Genesis environment containing the robot + command: The commanded XYZ position and Euler angles (in world frame), shape (num_envs, 6) where first 3 are position and last 3 are Euler angles. + pose_cmd_manager: The Pose command manager + pos_sensitivity: A lower value means the reward is more sensitive to position error + euler_sensitivity: A lower value means the reward is more sensitive to Euler angle error + entity_attr: The attribute name of the entity in the environment. + entity_manager: The entity manager for the entity. + + Returns: + torch.Tensor: Penalty for base position and Euler angles away from target + """ + assert ( + command is not None or pose_cmd_manager is not None + ), "Either command or pose_cmd_manager must be provided to command_tracking_pose" + + if entity_manager is not None: + entity = entity_manager.entity + else: + entity = getattr(env, entity_attr) + + if link_name is None or link_name == "": + link = entity.base_link + else: + try: + link = entity.get_link(link_name) + except Exception as e: + raise ValueError( + f"Link name '{link_name}' not found in entity '{entity_attr}'." + ) + + link_pos = link.get_pos() + link_euler = quat_to_xyz(link.get_quat()) + + if pose_cmd_manager is not None: + command = pose_cmd_manager.command + command_pos = command[:, :3] + command_euler = command[:, 3:6] + + pos_error = torch.sum(torch.square(command_pos - link_pos), dim=1) + + euler_error = torch.sum( + torch.square( + torch.min( + torch.abs(command_euler - link_euler), + 2 * torch.pi - torch.abs(command_euler - link_euler), + ) + ), + dim=1, + ) + + pos_penalty = torch.exp(-pos_error / pos_sensitivity) + euler_penalty = torch.exp(-euler_error / euler_sensitivity) + + total_penalty = pos_penalty + euler_penalty + + return total_penalty + + """ Velocity Command Rewards """