diff --git a/.gitignore b/.gitignore index 43d874f..804e65a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +.vscode/ # Byte-compiled / optimized / DLL files __pycache__/ *.py[codz] diff --git a/examples/rough_terrain/environment.py b/examples/rough_terrain/environment.py index fda9d24..6270625 100644 --- a/examples/rough_terrain/environment.py +++ b/examples/rough_terrain/environment.py @@ -13,6 +13,7 @@ VelocityCommandManager, TerrainManager, ContactManager, + ActuatorManager ) from genesis_forge.mdp import reset, rewards, terminations diff --git a/examples/skrl/README.md b/examples/skrl/README.md new file mode 100644 index 0000000..a8a9787 --- /dev/null +++ b/examples/skrl/README.md @@ -0,0 +1,40 @@ +# Go2 Simple Locomotion Example + +A simple program that teaches the Go2 robot to walk forward. + +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-walking/videos`. + +https://github.com/user-attachments/assets/be46df1b-35e5-4b5b-9bbc-f543210dd463 + + +## Evaluation + +Now you can view the trained policy: + +```bash +python ./eval.py ./logs/go2-walking/ +``` diff --git a/examples/skrl/environment.py b/examples/skrl/environment.py new file mode 100644 index 0000000..32caf10 --- /dev/null +++ b/examples/skrl/environment.py @@ -0,0 +1,272 @@ +""" +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, +) +from genesis_forge.mdp import reset, rewards, terminations +try: + import rclpy + from rclpy.node import Node +except ImportError: + pass + + +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 Go2SimpleEnv(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, + deploy_with_ros: bool = False, + ros_node = None, + ): + super().__init__( + num_envs=num_envs, + dt=dt, + max_episode_length_sec=max_episode_length_s, + max_episode_random_scaling=0.1, + deploy_with_ros=deploy_with_ros + ) + self._deploy_with_ros = deploy_with_ros + self._ros_node = ros_node + # print("deploy_with_ros_is:",self._deploy_with_ros) + # exit(0) + # Set the commanded robot direction to be 0.5 along the X axis, for all environments + self.target_command = torch.zeros( + (self.num_envs, 3), device=gs.device, dtype=gs.tc_float + ) + self.target_command[:, 0] = ( + TARGET_X_VELOCITY # Linear velocity along the X axis + ) + + # Construct the scene + if not self._deploy_with_ros: + 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", + ros_node=self._ros_node if self._deploy_with_ros and self.scene.num_envs==1 else None, + on_reset={ + # Reset the robot's initial position + "position": { + "fn": reset.position, + "params": { + "position": INITIAL_BODY_POSITION, + "quat": INITIAL_QUAT, + "zero_velocity": True, + }, + }, + }, + ) + + ## + # 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, + ros_node=self._ros_node if self._deploy_with_ros and self.scene.num_envs==1 else None, + ) + 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={ + "base_height_target": { + "weight": -50.0, + "fn": rewards.base_height, + "params": { + "target_height": 0.3, + "entity_attr": "robot", + }, + }, + "tracking_lin_vel": { + "weight": 1.0, + "fn": rewards.command_tracking_lin_vel, + "params": { + "command": self.target_command[:, :2], + "entity_manager": self.robot_manager, + }, + }, + "tracking_ang_vel": { + "weight": 0.2, + "fn": rewards.command_tracking_ang_vel, + "params": { + "commanded_ang_vel": self.target_command[:, 2], + "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.actuator_manager.get_dofs_position(), + }, + "dof_velocity": { + "fn": lambda env: self.actuator_manager.get_dofs_velocity(), + "scale": 0.05, + }, + "actions": { + "fn": lambda env: self.action_manager.get_actions(), + }, + }, + ) + + def build(self): + super().build(delpy_with_ros=self._deploy_with_ros) + self.camera.follow_entity(self.robot) + + # def step(self): + # super().step() + # if rclpy.ok() and self._deploy_with_ros: + # rclpy.spin_once(self._ros_node,timeout_sec=0) + \ No newline at end of file diff --git a/examples/skrl/eval.py b/examples/skrl/eval.py new file mode 100644 index 0000000..12f341d --- /dev/null +++ b/examples/skrl/eval.py @@ -0,0 +1,87 @@ +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 Go2SimpleEnv + +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-simple" + +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) +parser.add_argument("-ros", "--deploy_with_ros", type=bool, default=False) +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 = Go2SimpleEnv(num_envs=1, headless=False,deploy_with_ros=args.deploy_with_ros) + 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/skrl/train.py b/examples/skrl/train.py new file mode 100644 index 0000000..d7538a0 --- /dev/null +++ b/examples/skrl/train.py @@ -0,0 +1,158 @@ +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, + SkrlEnvWapper, +) +from environment import Go2SimpleEnv + +import torch.nn as nn + +from skrl.agents.torch.ppo import PPO, PPO_DEFAULT_CONFIG +from skrl.utils.runner.torch import Runner +from skrl.memories.torch import RandomMemory +from skrl.models.torch import DeterministicMixin, GaussianMixin, Model +from skrl.resources.preprocessors.torch import RunningStandardScaler +from skrl.resources.schedulers.torch import KLAdaptiveLR +from skrl.trainers.torch import SequentialTrainer + +class Shared(GaussianMixin, DeterministicMixin, Model): + def __init__(self, observation_space, action_space, device, clip_actions=False, + clip_log_std=True, min_log_std=-20, max_log_std=2, reduction="sum"): + Model.__init__(self, observation_space, action_space, device) + GaussianMixin.__init__(self, clip_actions, clip_log_std, min_log_std, max_log_std, reduction) + DeterministicMixin.__init__(self, clip_actions) + + print("num_observations:",self.num_observations) + print("num_actions:",self.num_actions) + # exit(0) + self.net = nn.Sequential(nn.Linear(self.num_observations, 512), + nn.ELU(), + nn.Linear(512, 256), + nn.ELU(), + nn.Linear(256, 128), + nn.ELU(), + nn.Linear(128, 64), + nn.ELU()) + + self.mean_layer = nn.Linear(64, self.num_actions) + self.log_std_parameter = nn.Parameter(torch.ones(self.num_actions)) + + self.value_layer = nn.Linear(64, 1) + + def act(self, inputs, role): + if role == "policy": + return GaussianMixin.act(self, inputs, role) + elif role == "value": + return DeterministicMixin.act(self, inputs, role) + + def compute(self, inputs, role): + # for key, value in inputs.items(): + # print(key,"!!:!!",value.shape) + if role == "policy": + self._shared_output = self.net(inputs["states"]) + return self.mean_layer(self._shared_output), self.log_std_parameter, {} + elif role == "value": + shared_output = self.net(inputs["states"]) if self._shared_output is None else self._shared_output + self._shared_output = None + return self.value_layer(shared_output), {} + +EXPERIMENT_NAME = "go2-not-simple" + +parser = argparse.ArgumentParser(add_help=True) +parser.add_argument("-n", "--num_envs", type=int, default=4096) +parser.add_argument("--max_iterations", type=int, default=1000) +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 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}") + + # Create environment + env = Go2SimpleEnv(num_envs=args.num_envs, headless=True,deploy_with_ros=False) + + # 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 = SkrlEnvWapper(env) + env.build() + env.reset() + + memory = RandomMemory(memory_size=60, num_envs=env.num_envs, device=gs.device) + + cfg = PPO_DEFAULT_CONFIG.copy() + cfg["rollouts"] = 60 # memory_size + cfg["learning_epochs"] = 5 + cfg["mini_batches"] = 4 # 96 * 4096 / 98304 + cfg["discount_factor"] = 0.99 + cfg["lambda"] = 0.95 + cfg["learning_rate"] = 1e-3 + cfg["learning_rate_scheduler"] = KLAdaptiveLR + cfg["learning_rate_scheduler_kwargs"] = {"kl_threshold": 0.01, "min_lr": 5e-4} + cfg["random_timesteps"] = 0 + cfg["learning_starts"] = 0 + cfg["grad_norm_clip"] = 1.0 + cfg["ratio_clip"] = 0.2 + cfg["value_clip"] = 0.2 + cfg["clip_predicted_values"] = True + cfg["entropy_loss_scale"] = 0.01 + cfg["value_loss_scale"] = 1.0 + cfg["kl_threshold"] = 0 + cfg["rewards_shaper"] = None + cfg["time_limit_bootstrap"] = True + cfg["state_preprocessor"] = RunningStandardScaler + cfg["state_preprocessor_kwargs"] = {"size": env.observation_space, "device": gs.device} + cfg["value_preprocessor"] = RunningStandardScaler + cfg["value_preprocessor_kwargs"] = {"size": 1, "device": gs.device} + # logging to TensorBoard and write checkpoints (in timesteps) + cfg["experiment"]["write_interval"] = 60 + cfg["experiment"]["checkpoint_interval"] = 100 + models = {} + models["policy"] = Shared(env.observation_space, env.action_space, gs.device) + models["value"] = models["policy"] # same instance: shared model + agent = PPO(models=models, + memory=memory, + cfg=cfg, + observation_space=env.observation_space, + action_space=env.action_space, + device=gs.device) + cfg_trainer = {"timesteps": args.max_iterations*cfg["learning_epochs"], "headless": True} + trainer = SequentialTrainer(cfg=cfg_trainer, env=env, agents=agent) + # Train + print("💪 Training model...") + trainer.train() + env.close() + + +if __name__ == "__main__": + main() diff --git a/examples/skrl_with_sensors/README.md b/examples/skrl_with_sensors/README.md new file mode 100644 index 0000000..a8a9787 --- /dev/null +++ b/examples/skrl_with_sensors/README.md @@ -0,0 +1,40 @@ +# Go2 Simple Locomotion Example + +A simple program that teaches the Go2 robot to walk forward. + +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-walking/videos`. + +https://github.com/user-attachments/assets/be46df1b-35e5-4b5b-9bbc-f543210dd463 + + +## Evaluation + +Now you can view the trained policy: + +```bash +python ./eval.py ./logs/go2-walking/ +``` diff --git a/examples/skrl_with_sensors/environment.py b/examples/skrl_with_sensors/environment.py new file mode 100644 index 0000000..d8f4407 --- /dev/null +++ b/examples/skrl_with_sensors/environment.py @@ -0,0 +1,280 @@ +""" +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, + CameraManager +) +from genesis_forge.mdp import reset, rewards, terminations +try: + import rclpy + from rclpy.node import Node +except ImportError: + pass + + +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 Go2SimpleEnv(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, + deploy_with_ros: bool = False, + ros_node = None, + ): + super().__init__( + num_envs=num_envs, + dt=dt, + max_episode_length_sec=max_episode_length_s, + max_episode_random_scaling=0.1, + ) + self._deploy_with_ros = deploy_with_ros + self._ros_node = ros_node + self.target_command = torch.zeros( + (self.num_envs, 3), device=gs.device, dtype=gs.tc_float + ) + self.target_command[:, 0] = ( + TARGET_X_VELOCITY # Linear velocity along the X axis + ) + + # Construct the scene + if not self._deploy_with_ros: + 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", + ros_node=self._ros_node if self._deploy_with_ros and self.scene.num_envs==1 else None, + on_reset={ + # Reset the robot's initial position + "position": { + "fn": reset.position, + "params": { + "position": INITIAL_BODY_POSITION, + "quat": INITIAL_QUAT, + "zero_velocity": True, + }, + }, + }, + ) + + ## + # 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, + ros_node=self._ros_node if self._deploy_with_ros and self.scene.num_envs==1 else None, + ) + self.action_manager = PositionActionManager( + self, + scale=0.25, + clip=(-100.0, 100.0), + use_default_offset=True, + actuator_manager=self.actuator_manager, + ) + self.camera_manager=CameraManager( + self, + sensor_name="front_cam", + link_name="front_camera_link", + delay=0.5, + read_frequency=12, + res=(256,256), + ) + ## + # Rewards + RewardManager( + self, + logging_enabled=True, + cfg={ + "base_height_target": { + "weight": -50.0, + "fn": rewards.base_height, + "params": { + "target_height": 0.3, + "entity_attr": "robot", + }, + }, + "tracking_lin_vel": { + "weight": 1.0, + "fn": rewards.command_tracking_lin_vel, + "params": { + "command": self.target_command[:, :2], + "entity_manager": self.robot_manager, + }, + }, + "tracking_ang_vel": { + "weight": 0.2, + "fn": rewards.command_tracking_ang_vel, + "params": { + "commanded_ang_vel": self.target_command[:, 2], + "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.actuator_manager.get_dofs_position(), + }, + "dof_velocity": { + "fn": lambda env: self.actuator_manager.get_dofs_velocity(), + "scale": 0.05, + }, + "actions": { + "fn": lambda env: self.action_manager.get_actions(), + }, + "front_image": { + "fn": lambda env: self.camera_manager.get_rgb_img(), + }, + }, + obs_to_discard=["front_image"] + ) + + def build(self): + super().build() + self.camera.follow_entity(self.robot) + + # def step(self): + # super().step() + # if rclpy.ok() and self._deploy_with_ros: + # rclpy.spin_once(self._ros_node,timeout_sec=0) + \ No newline at end of file diff --git a/examples/skrl_with_sensors/eval.py b/examples/skrl_with_sensors/eval.py new file mode 100644 index 0000000..12f341d --- /dev/null +++ b/examples/skrl_with_sensors/eval.py @@ -0,0 +1,87 @@ +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 Go2SimpleEnv + +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-simple" + +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) +parser.add_argument("-ros", "--deploy_with_ros", type=bool, default=False) +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 = Go2SimpleEnv(num_envs=1, headless=False,deploy_with_ros=args.deploy_with_ros) + 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/skrl_with_sensors/train.py b/examples/skrl_with_sensors/train.py new file mode 100644 index 0000000..966d119 --- /dev/null +++ b/examples/skrl_with_sensors/train.py @@ -0,0 +1,196 @@ +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, + SkrlEnvWapper, +) +from environment import Go2SimpleEnv + +import torch.nn as nn + +from skrl.agents.torch.ppo import PPO, PPO_DEFAULT_CONFIG +from skrl.utils.runner.torch import Runner +from skrl.memories.torch import RandomMemory +from skrl.models.torch import DeterministicMixin, GaussianMixin, Model +from skrl.resources.preprocessors.torch import RunningStandardScaler +from skrl.resources.schedulers.torch import KLAdaptiveLR +from skrl.trainers.torch import SequentialTrainer + + +import torch +import torch.nn as nn + +class Shared(GaussianMixin, DeterministicMixin, Model): + def __init__(self, observation_space, action_space, device, clip_actions=False, + clip_log_std=True, min_log_std=-20, max_log_std=2, reduction="sum"): + Model.__init__(self, observation_space, action_space, device) + GaussianMixin.__init__(self, clip_actions, clip_log_std, min_log_std, max_log_std, reduction) + DeterministicMixin.__init__(self, clip_actions) + # print("obs space:",self.observation_space) + # exit(0) + self.conv_layers = nn.Sequential( + nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1), + nn.ReLU(), + nn.MaxPool2d(kernel_size=2, stride=2), + nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1), + nn.ReLU(), + nn.MaxPool2d(kernel_size=2, stride=2), + nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), + nn.ReLU(), + nn.MaxPool2d(kernel_size=2, stride=2) + ) + + self.net = nn.Sequential( + nn.Linear(109, 64), + nn.ELU() + ) + + self.mean_layer = nn.Linear(64, self.num_actions) + self.log_std_parameter = nn.Parameter(torch.ones(self.num_actions)) + self.value_layer = nn.Linear(64, 1) + self._shared_output = None + + def act(self, inputs, role): + if role == "policy": + return GaussianMixin.act(self, inputs, role) + elif role == "value": + return DeterministicMixin.act(self, inputs, role) + + def compute(self, inputs, role): + if role == "policy": + states = inputs["states"] + # print("obs space:",self.observation_space) + # print("states:",states) + # exit(0) + space = self.tensor_to_space(states, self.observation_space) + # print("front_image_shape",space.shape) + image_features = self.conv_layers(space["front_image"].permute(0, 3, 1, 2).float()) + # Apply global average pooling across spatial dimensions (height and width) + image_features = image_features.mean(dim=[2, 3]) # Global average pooling (across height and width) + # Flatten the tensor to have a single dimension of features + image_features = torch.flatten(image_features, 1) + proprio_features = torch.cat([ + space[key].view(states.shape[0], -1) for key in space.keys() if key != "front_image" + ], dim=1) + + self._shared_output = self.net(torch.cat([image_features, proprio_features], dim=1)) + + return self.mean_layer(self._shared_output), self.log_std_parameter, {} + + elif role == "value": + if self._shared_output is None: + states = inputs["states"] + space = self.tensor_to_space(states, self.observation_space) + image_features = self.conv_layers(space["front_image"].permute(0, 3, 1, 2).float()) + # Apply global average pooling across spatial dimensions (height and width) + image_features = image_features.mean(dim=[2, 3]) # Global average pooling (across height and width) + # Flatten the tensor to have a single dimension of features + image_features = torch.flatten(image_features, 1) + proprio_features = torch.cat([ + space[key].view(states.shape[0], -1) for key in space.keys() if key != "front_image" + ], dim=1) + + self._shared_output = self.net(torch.cat([image_features, proprio_features], dim=1)) + + return self.value_layer(self._shared_output), {} + + +EXPERIMENT_NAME = "go2-camera" + +parser = argparse.ArgumentParser(add_help=True) +parser.add_argument("-n", "--num_envs", type=int, default=25) +parser.add_argument("--max_iterations", type=int, default=1000) +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 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}") + + # Create environment + env = Go2SimpleEnv(num_envs=args.num_envs, headless=True,deploy_with_ros=False) + + # 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 = SkrlEnvWapper(env) + env.build() + env.reset() + + memory = RandomMemory(memory_size=4, num_envs=env.num_envs, device=gs.device) + + cfg = PPO_DEFAULT_CONFIG.copy() + cfg["rollouts"] = 4 # memory_size + cfg["learning_epochs"] = 5 + cfg["mini_batches"] = 2 # 96 * 4096 / 98304 + cfg["discount_factor"] = 0.99 + cfg["lambda"] = 0.95 + cfg["learning_rate"] = 1e-3 + cfg["learning_rate_scheduler"] = KLAdaptiveLR + cfg["learning_rate_scheduler_kwargs"] = {"kl_threshold": 0.01, "min_lr": 5e-4} + cfg["random_timesteps"] = 0 + cfg["learning_starts"] = 0 + cfg["grad_norm_clip"] = 1.0 + cfg["ratio_clip"] = 0.2 + cfg["value_clip"] = 0.2 + cfg["clip_predicted_values"] = True + cfg["entropy_loss_scale"] = 0.01 + cfg["value_loss_scale"] = 1.0 + cfg["kl_threshold"] = 0 + cfg["rewards_shaper"] = None + cfg["time_limit_bootstrap"] = True + cfg["state_preprocessor"] = RunningStandardScaler + cfg["state_preprocessor_kwargs"] = {"size": env.observation_space, "device": gs.device} + cfg["value_preprocessor"] = RunningStandardScaler + cfg["value_preprocessor_kwargs"] = {"size": 1, "device": gs.device} + # logging to TensorBoard and write checkpoints (in timesteps) + cfg["experiment"]["write_interval"] = 60 + cfg["experiment"]["checkpoint_interval"] = 100 + models = {} + models["policy"] = Shared(env.observation_space, env.action_space, gs.device) + models["value"] = models["policy"] # same instance: shared model + agent = PPO(models=models, + memory=memory, + cfg=cfg, + observation_space=env.observation_space, + action_space=env.action_space, + device=gs.device) + cfg_trainer = {"timesteps": args.max_iterations*cfg["learning_epochs"], "headless": True} + trainer = SequentialTrainer(cfg=cfg_trainer, env=env, agents=agent) + # Train + print("💪 Training model...") + trainer.train() + env.close() + + +if __name__ == "__main__": + main() diff --git a/genesis_forge/__init__.py b/genesis_forge/__init__.py index 3c73d74..a24a9c5 100644 --- a/genesis_forge/__init__.py +++ b/genesis_forge/__init__.py @@ -1,8 +1,12 @@ from .genesis_env import GenesisEnv, EnvMode +from .ros2_env import Ros2Env from .managed_env import ManagedEnvironment +from .ros2_managed_env import Ros2ManagedEnvironment __all__ = [ "GenesisEnv", + "Ros2Env", "ManagedEnvironment", + "Ros2ManagedEnvironment", "EnvMode", ] diff --git a/genesis_forge/managed_env.py b/genesis_forge/managed_env.py index 1d9d15d..577fdde 100644 --- a/genesis_forge/managed_env.py +++ b/genesis_forge/managed_env.py @@ -15,6 +15,11 @@ RewardManager, TerminationManager, ActuatorManager, + CameraManager, + ImuManager, + DepthCameraManager, + SphericalRaycasterManager, + GridRaycasterManager ) @@ -22,6 +27,11 @@ class ManagersDict(TypedDict): actuator: ActuatorManager | None contact: list[ContactManager] entity: list[EntityManager] + imu_sensor: list[ImuManager] + camera_sensor: list[CameraManager] + depth_camera_sensor: list[DepthCameraManager] + grid_raycaster_sensor: list[GridRaycasterManager] + spherical_raycaster_sensor: list[SphericalRaycasterManager] command: list[CommandManager] terrain: list[TerrainManager] action: PositionActionManager | None @@ -127,12 +137,16 @@ def __init__( max_episode_random_scaling=max_episode_random_scaling, extras_logging_key=extras_logging_key, ) - self.managers: ManagersDict = { "contact": [], "entity": [], "command": [], "terrain": [], + "imu_sensor":[], + "camera_sensor":[], + "depth_camera_sensor":[], + "grid_raycaster_sensor":[], + "spherical_raycaster_sensor":[], # there can only be one of each of these "actuator": None, "action": None, @@ -254,8 +268,8 @@ def build(self): Builds the environment before the first step. The Genesis scene and all the scene entities must be added before calling this method. """ - super().build() self.config() + super().build() for terrain_manager in self.managers["terrain"]: terrain_manager.build() @@ -263,8 +277,25 @@ def build(self): self.managers["actuator"].build() if self.managers["action"] is not None: self.managers["action"].build() + for contact_manager in self.managers["contact"]: contact_manager.build() + + for imu_manager in self.managers["imu_sensor"]: + imu_manager.build() + + for camera_manager in self.managers["camera_sensor"]: + camera_manager.build() + + for depth_camera_manager in self.managers["depth_camera_sensor"]: + depth_camera_manager.build() + + for grid_raycaster_manager in self.managers["grid_raycaster_sensor"]: + grid_raycaster_manager.build() + + for spherical_raycaster_manager in self.managers["spherical_raycaster_sensor"]: + spherical_raycaster_manager.build() + if self.managers["termination"] is not None: self.managers["termination"].build() if self.managers["reward"] is not None: @@ -277,7 +308,7 @@ def build(self): obs.build() def step( - self, actions: torch.Tensor + self, actions: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: """ Performs a step in all environments with the given actions. @@ -294,6 +325,7 @@ def step( # Execute the actions and a simulation step if self.managers["action"] is not None: self.managers["action"].step(actions) + self.scene.step() # Update entity managers @@ -303,6 +335,21 @@ def step( # Calculate contact forces for contact_manager in self.managers["contact"]: contact_manager.step() + + for imu_manager in self.managers["imu_sensor"]: + imu_manager.step() + + for camera_manager in self.managers["camera_sensor"]: + camera_manager.step() + + for depth_camera_manager in self.managers["depth_camera_sensor"]: + depth_camera_manager.step() + + for grid_raycaster_manager in self.managers["grid_raycaster_sensor"]: + grid_raycaster_manager.step() + + for spherical_raycaster_manager in self.managers["spherical_raycaster_sensor"]: + spherical_raycaster_manager.step() # Calculate termination and truncation reset_env_idx = None @@ -329,7 +376,6 @@ def step( # Get observations obs = self.get_observations() - return ( obs, rewards, diff --git a/genesis_forge/managers/__init__.py b/genesis_forge/managers/__init__.py index db47035..91a6b46 100644 --- a/genesis_forge/managers/__init__.py +++ b/genesis_forge/managers/__init__.py @@ -3,8 +3,17 @@ 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 .contact import ContactManager +from .command import ( + CommandManager, + PositionCommandManager, + PoseCommandManager, + VelocityCommandManager) +from .sensors.contact.contact_manager import ContactManager +from .sensors.imu_manager import ImuManager +from .sensors.camera_manager import CameraManager +from .sensors.depth_camera_manager import DepthCameraManager +from .sensors.grid_raycaster_manager import GridRaycasterManager +from .sensors.spherical_raycaster_manager import SphericalRaycasterManager from .terrain_manager import TerrainManager from .entity_manager import EntityManager from .observation_manager import ObservationManager @@ -19,10 +28,17 @@ "RewardManager", "TerminationManager", "CommandManager", + "PositionCommandManager", + "PoseCommandManager", "VelocityCommandManager", "PositionActionManager", "PositionWithinLimitsActionManager", "ContactManager", + "ImuManager", + "CameraManager", + "DepthCameraManager", + "GridRaycasterManager", + "SphericalRaycasterManager", "TerrainManager", "EntityManager", "ObservationManager", diff --git a/genesis_forge/managers/action/__init__.py b/genesis_forge/managers/action/__init__.py index 7fa82a2..c9a4889 100644 --- a/genesis_forge/managers/action/__init__.py +++ b/genesis_forge/managers/action/__init__.py @@ -1,9 +1,17 @@ from .base import BaseActionManager from .position_action_manager import PositionActionManager from .position_within_limits import PositionWithinLimitsActionManager +from .velocity_action_manager import VelocityActionManager +from .force_action_manager import ForceActionManager +from .force_within_limits import ForceWithinLimitsActionManager +from .hybrid_action_manager import HybridActionManager __all__ = [ "BaseActionManager", "PositionActionManager", "PositionWithinLimitsActionManager", + "VelocityActionManager", + "ForceActionManager", + "ForceWithinLimitsActionManager", + "HybridActionManager" ] diff --git a/genesis_forge/managers/action/force_action_manager.py b/genesis_forge/managers/action/force_action_manager.py new file mode 100644 index 0000000..a59c391 --- /dev/null +++ b/genesis_forge/managers/action/force_action_manager.py @@ -0,0 +1,360 @@ +from __future__ import annotations +import re +import torch +import genesis as gs +import numpy as np +from gymnasium import spaces +from typing import Any, Callable, TypeVar +from deprecated import deprecated +from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.managers.action.base import BaseActionManager +from genesis_forge.values import ensure_dof_pattern +from genesis_forge.managers.actuator import ActuatorManager + +deprecated_arg_names = [ + "joint_names", + "default_pos", + "pd_kp", + "pd_kv", + "max_force", + "damping", + "stiffness", + "frictionloss", + "noise_scale", +] + +T = TypeVar("T") + + +class ForceActionManager(BaseActionManager): + """ + Converts actions to DOF forces, using affine transformations (scale). + + .. math:: + + force = scaling * action + + Args: + env: The environment to manage the DOF actuators for. + actuator_manager: The actuator manager which is used to setup and control the DOF joints. + scale: How much to scale the action. + clip: Clip the action values to the range. If omitted, the action values will automatically be clipped to the joint limits. + quiet_action_errors: Whether to quiet action errors. + delay_step: The number of steps to delay the actions for. + This is an easy way to emulate the latency in the system. + + Example:: + + class MyEnv(ManagedEnvironment): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # ...define scene and robot... + + def config(self): + self.actuator_manager = ActuatorManager( + self, + joint_names=".*", + default_pos={".*": 0.0}, + kp=50, + kv=0.5, + max_force=8.0, + ) + self.action_manager = ForceActionManager( + self, + scale=0.5, + actuator_manager=self.actuator_manager, + ) + + Example using the manager directly:: + + class MyEnv(GenesisEnv): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # ...define scene and robot... + + self.actuator_manager = ActuatorManager( + self, + joint_names=".*", + default_pos={".*": 0.0}, + kp=50, + kv=0.5, + max_force=8.0, + ) + self.action_manager = ForceActionManager( + self, + scale=0.5, + ) + + def build(self): + super().build() + self.actuator_manager.build() + self.action_manager.build() + + step(self, actions: torch.Tensor) -> None: + super().step(actions) + self.action_manager.step(actions) + + # ...do other step things... + + reset(self, envs_idx: list[int] = None) -> None: + super().reset(envs_idx) + self.actuator_manager.reset(envs_idx) + self.action_manager.reset(envs_idx) + + # ...do other reset things... + + + """ + + def __init__( + self, + env: GenesisEnv, + actuator_manager: ActuatorManager | None = None, + scale: float | dict[str, float] = 1.0, + clip: tuple[float, float] | dict[str, tuple[float, float]] = None, + action_handler: Callable[[torch.Tensor], None] = None, + quiet_action_errors: bool = False, + delay_step: int = 0, + **kwargs, + ): + super().__init__(env, delay_step) + self._scale_cfg = ensure_dof_pattern(scale) + self._clip_cfg = ensure_dof_pattern(clip) + self._quiet_action_errors = quiet_action_errors + self._enabled_dof = None + self._actuator_manager = actuator_manager + + self._dofs_pos_buffer: torch.Tensor = None + + # Deprecated actuator parameters + deprecated_actuator_args = { + key: kwargs[key] for key in deprecated_arg_names if key in kwargs + } + if len(deprecated_actuator_args) > 0: + dep_list = ", ".join(deprecated_actuator_args.keys()) + if self._actuator_manager is not None: + raise ValueError( + f"Cannot set both actuator_manager and deprecated actuator parameters: {dep_list}" + ) + print( + f"Actuator arguments are deprecated in the action manager, instead define an ActuatorManager ({dep_list})" + ) + self._actuator_manager = ActuatorManager( + env, + joint_names=kwargs.get("joint_names", ".*"), + default_pos=kwargs.get("default_pos", {".*": 0.0}), + kp=kwargs.get("pd_kp", None), + kv=kwargs.get("pd_kv", None), + max_force=kwargs.get("max_force", None), + damping=kwargs.get("damping", None), + stiffness=kwargs.get("stiffness", None), + frictionloss=kwargs.get("frictionloss", None), + default_noise_scale=kwargs.get("noise_scale", 0.0), + ) + if self._actuator_manager is None: + raise ValueError("No ActuatorManager provided.") + + """ + Properties + """ + + @property + def actuators(self) -> ActuatorManager: + """ + Get the actuator manager. + """ + return self._actuator_manager + + @property + def num_actions(self) -> int: + """ + Get the number of actions. + """ + return self._actuator_manager.num_dofs + + @property + def action_space(self) -> tuple[float, float]: + """ + Returns the actions space for the environment, based on the number of DOFs defined in this action manager. + """ + return spaces.Box( + low=-np.inf, + high=np.inf, + shape=(self.num_actions,), + dtype=np.float32, + ) + + @property + def dofs_idx(self) -> list[int]: + """ + Get the indices of the DOF that are enabled (via joint_names). + """ + return self._actuator_manager.dofs_idx + + @property + def default_dofs_pos(self) -> torch.Tensor: + """ + Return the default DOF positions. + """ + return self._actuator_manager.default_dofs_pos + + """ + DOF Getters + """ + + @deprecated( + version="0.3,0", + reason="Use the actuator manager directly.", + ) + def get_dofs_position(self, noise: float = 0.0): + """ + Deprecated: Use the actuator manager directly. + + Return the current position of the enabled DOFs. + This is a wrapper for `RigidEntity.get_dofs_position`. + + Args: + noise: The maximum amount of random noise to add to the position values returned. + """ + return self._actuator_manager.get_dofs_position(noise) + + @deprecated( + version="0.3,0", + reason="Use the actuator manager directly.", + ) + def get_dofs_velocity(self, noise: float = 0.0, clip: tuple[float, float] = None): + """ + Deprecated: Use the actuator manager directly. + + Return the current velocity of the enabled DOFs. + This is a wrapper for `RigidEntity.get_dofs_velocity`. + + Args: + noise: The maximum amount of random noise to add to the velocity values returned. + clip: Clip the velocity returned. + """ + return self._actuator_manager.get_dofs_velocity(noise, clip) + + @deprecated( + version="0.3,0", + reason="Use the actuator manager directly.", + ) + def get_dofs_force(self, noise: float = 0.0, clip_to_max_force: bool = False): + """ + Deprecated: Use the actuator manager directly. + + Return the force experienced by the enabled DOFs. + This is a wrapper for `RigidEntity.get_dofs_force`. + + Args: + noise: The maximum amount of random noise to add to the force values returned. + clip_to_max_force: Clip the force returned to the maximum force defined by the `max_force` parameter. + + Returns: + The force experienced by the enabled DOFs. + """ + return self._actuator_manager.get_dofs_force(noise, clip_to_max_force) + + """ + Operations + """ + + def build(self): + """ + Builds the manager and initialized all the buffers. + """ + + # Define the clip values + lower_limit, upper_limit = self._actuator_manager.get_dofs_limits() + self._clip_values = torch.stack([lower_limit, upper_limit], dim=1) + if self._clip_cfg is not None: + self._get_dof_value_tensor(self._clip_cfg, output=self._clip_values) + + # Scale + self._scale_values = None + if self._scale_cfg is not None: + self._scale_values = self._get_dof_value_tensor(self._scale_cfg) + + def step(self, actions: torch.Tensor) -> torch.Tensor: + """ + Take the incoming actions for this step and handle them. + + Args: + actions: The incoming step actions to handle. + """ + if not self.enabled: + return + actions = super().step(actions) + self._actions = self.handle_actions(actions) + return self._actions + + def handle_actions(self, actions: torch.Tensor) -> torch.Tensor: + """ + Converts the actions to position commands, and send them to the DOF actuators. + Override this function if you want to change the action handling logic. + + Args: + actions: The incoming step actions to handle. + + Returns: + The processed and handled actions. + """ + + # Validate actions + if not self._quiet_action_errors: + if torch.isnan(actions).any(): + print(f"ERROR: NaN actions received! Actions: {actions}") + if torch.isinf(actions).any(): + print(f"ERROR: Infinite actions received! Actions: {actions}") + + # Process actions + actions = actions * self._scale_values + actions = torch.clamp( + actions, + min=self._clip_values[:, 0], + max=self._clip_values[:, 1], + ) + + # Set target positions + self._actuator_manager.control_dofs_force(actions) + + return actions + + """ + Internal methods + """ + + def _get_dof_value_tensor( + self, + values: float | dict, + default_value: T = 0.0, + output: torch.Tensor | list[Any] | None = None, + ) -> torch.Tensor: + """ + Given a DofValue dict, loop over the entries, and set the value to the DOF indices (from the actuator) that match the pattern. + + Args: + values: The DOF value to convert (for example: `{".*": 50}`). + + Returns: + A list of values for the DOF indices. + For example, for 4 DOFs: [50, 50, 50, 50] + """ + is_set = [False] * self.num_actions + dof_names = self._actuator_manager.dofs_names + if output is None: + output = torch.zeros( + self.num_actions, device=gs.device, dtype=gs.tc_float + ).fill_(default_value) + for pattern, value in values.items(): + found = False + for i, name in enumerate[str](dof_names): + if not is_set[i] and re.match(f"^{pattern}$", name): + if isinstance(value, (list, tuple)): + value = torch.tensor(value, device=gs.device) + is_set[i] = True + output[i] = value + found = True + if not found: + raise RuntimeError(f"Joint DOF '{pattern}' not found.") + return output diff --git a/genesis_forge/managers/action/force_within_limits.py b/genesis_forge/managers/action/force_within_limits.py new file mode 100644 index 0000000..65f8e29 --- /dev/null +++ b/genesis_forge/managers/action/force_within_limits.py @@ -0,0 +1,102 @@ +from __future__ import annotations +import torch +from typing import Callable + +from genesis_forge.genesis_env import GenesisEnv +from .force_action_manager import ForceActionManager +from genesis_forge.managers.actuator import ActuatorManager + + +class ForceWithinLimitsActionManager(ForceActionManager): + """ + This is similar to `ForceActionManager` but converts actions from the range -1.0 - 1.0 to DOF force within the limits of the actuators. + + Args: + env: The environment to manage the DOF actuators for. + actuator_manager: The actuator manager which is used to setup and control the DOF joints. + action_handler: A function to handle the actions. + quiet_action_errors: Whether to quiet action errors. + delay_step: The number of steps to delay the actions for. + This is an easy way to emulate the latency in the system. + + Example:: + + class MyEnv(ManagedEnvironment): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def config(self): + self.actuator_manager = ActuatorManager( + self, + joint_names=".*", + default_pos={ + # Hip joints + "Leg[1-2]_Hip": -1.0, + "Leg[3-4]_Hip": 1.0, + # Femur joints + "Leg[1-4]_Femur": 0.5, + # Tibia joints + "Leg[1-4]_Tibia": 0.6, + }, + kp={".*": 50}, + kv={".*": 0.5}, + max_force={".*": 8.0}, + ) + self.action_manager = PositionalActionManager( + self, + actuator_manager=self.actuator_manager, + ) + + """ + + def __init__( + self, + env: GenesisEnv, + actuator_manager: ActuatorManager | None = None, + action_handler: Callable[[torch.Tensor], None] = None, + quiet_action_errors: bool = False, + delay_step: int = 0, + **kwargs, + ): + super().__init__( + env, + action_handler=action_handler, + quiet_action_errors=quiet_action_errors, + delay_step=delay_step, + actuator_manager=actuator_manager, + **kwargs, + ) + + """ + Lifecycle Operations + """ + + def build(self): + """ + Builds the manager and initialized all the buffers. + """ + lower, upper = self._actuator_manager.get_dofs_force_range() + lower = lower.unsqueeze(0).expand(self.env.num_envs, -1) + upper = upper.unsqueeze(0).expand(self.env.num_envs, -1) + self._offset = (upper + lower) * 0.5 + self._scale = (upper - lower) * 0.5 + + def handle_actions(self, actions: torch.Tensor) -> torch.Tensor: + """ + Converts the actions to force commands, and send them to the DOF actuators. + Override this function if you want to change the action handling logic. + + Args: + actions: The incoming step actions to handle. + + Returns: + The processed and handled actions. + """ + # Convert the action from -1 to 1, to absolute position within the actuator limits + actions.clamp_(-1.0, 1.0) + self._actions = actions * self._scale + + # Set target positions + self._actuator_manager.control_dofs_force(self._actions) + + return self._actions diff --git a/genesis_forge/managers/action/hybrid_action_manager.py b/genesis_forge/managers/action/hybrid_action_manager.py new file mode 100644 index 0000000..d55c5a6 --- /dev/null +++ b/genesis_forge/managers/action/hybrid_action_manager.py @@ -0,0 +1,436 @@ +from __future__ import annotations +import re +import torch +import genesis as gs +import numpy as np +from gymnasium import spaces +from typing import Any, Callable, TypeVar +from deprecated import deprecated +from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.managers.action.base import BaseActionManager +from genesis_forge.values import ensure_dof_pattern +from genesis_forge.managers.actuator import ActuatorManager + +deprecated_arg_names = [ + "joint_names", + "default_pos", + "pd_kp", + "pd_kv", + "max_force", + "damping", + "stiffness", + "frictionloss", + "noise_scale", +] + +T = TypeVar("T") + + +class HybridActionManager(BaseActionManager): + """ + Converts actions to DOF positions,velocities or forces based on the actuation type, using affine transformations (scale and offset). + + .. math:: + + position = offset + scaling * action + velocity = scaling * action + force = scaling * action + + If `use_default_offset` is `True`, the `offset` will be set to the `default_pos` value for each DOF/joint. + + Args: + env: The environment to manage the DOF actuators for. + actuator_manager: The actuator manager which is used to setup and control the DOF joints. + scale: How much to scale the action. + offset: Offset factor for the action. + use_default_offset: Whether to use default joint positions configured in the articulation asset as offset. Defaults to True. + clip: Clip the action values to the range. If omitted, the action values will automatically be clipped to the joint limits. + quiet_action_errors: Whether to quiet action errors. + delay_step: The number of steps to delay the actions for. + This is an easy way to emulate the latency in the system. + + Example:: + + class MyEnv(ManagedEnvironment): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # ...define scene and robot... + + def config(self): + self.actuator_manager = ActuatorManager( + self, + joint_names=".*", + default_pos={".*": 0.0}, + kp=50, + kv=0.5, + max_force=8.0, + ) + self.action_manager = HybridActionManager( + self, + scale=0.5, + use_default_offset=True, + actuator_manager=self.actuator_manager, + ) + + Example using the manager directly:: + + class MyEnv(GenesisEnv): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # ...define scene and robot... + + self.actuator_manager = ActuatorManager( + self, + joint_names=".*", + default_pos={".*": 0.0}, + kp=50, + kv=0.5, + max_force=8.0, + ) + self.action_manager = HybridActionManager( + self, + scale=0.5, + offset=0.0, + use_default_offset=True, + ) + + def build(self): + super().build() + self.actuator_manager.build() + self.action_manager.build() + + step(self, actions: torch.Tensor) -> None: + super().step(actions) + self.action_manager.step(actions) + + # ...do other step things... + + reset(self, envs_idx: list[int] = None) -> None: + super().reset(envs_idx) + self.actuator_manager.reset(envs_idx) + self.action_manager.reset(envs_idx) + + # ...do other reset things... + + + """ + + def __init__( + self, + env: GenesisEnv, + actuator_manager: ActuatorManager | None = None, + scale: float | dict[str, float] = 1.0, + offset: float | dict[str, float] = 0.0, + clip: tuple[float, float] | dict[str, tuple[float, float]] = None, + use_default_offset: bool = True, + action_handler: Callable[[torch.Tensor], None] = None, + quiet_action_errors: bool = False, + delay_step: int = 0, + **kwargs, + ): + super().__init__(env, delay_step) + self._offset_cfg = ensure_dof_pattern(offset) + self._scale_cfg = ensure_dof_pattern(scale) + self._clip_cfg = ensure_dof_pattern(clip) + self._quiet_action_errors = quiet_action_errors + self._enabled_dof = None + self._use_default_offset = use_default_offset + self._actuator_manager = actuator_manager + + self._dofs_pos_buffer: torch.Tensor = None + + if use_default_offset and offset != 0.0: + raise ValueError("Cannot set both use_default_offset and offset") + + # Deprecated actuator parameters + deprecated_actuator_args = { + key: kwargs[key] for key in deprecated_arg_names if key in kwargs + } + if len(deprecated_actuator_args) > 0: + dep_list = ", ".join(deprecated_actuator_args.keys()) + if self._actuator_manager is not None: + raise ValueError( + f"Cannot set both actuator_manager and deprecated actuator parameters: {dep_list}" + ) + print( + f"Actuator arguments are deprecated in the action manager, instead define an ActuatorManager ({dep_list})" + ) + self._actuator_manager = ActuatorManager( + env, + joint_names=kwargs.get("joint_names", ".*"), + default_pos=kwargs.get("default_pos", {".*": 0.0}), + kp=kwargs.get("pd_kp", None), + kv=kwargs.get("pd_kv", None), + max_force=kwargs.get("max_force", None), + damping=kwargs.get("damping", None), + stiffness=kwargs.get("stiffness", None), + frictionloss=kwargs.get("frictionloss", None), + default_noise_scale=kwargs.get("noise_scale", 0.0), + ) + if self._actuator_manager is None: + raise ValueError("No ActuatorManager provided.") + + """ + Properties + """ + + @property + def actuators(self) -> ActuatorManager: + """ + Get the actuator manager. + """ + return self._actuator_manager + + @property + def num_actions(self) -> int: + """ + Get the number of actions. + """ + return self._actuator_manager.num_dofs + + @property + def action_space(self) -> tuple[float, float]: + """ + Returns the actions space for the environment, based on the number of DOFs defined in this action manager. + """ + return spaces.Box( + low=-np.inf, + high=np.inf, + shape=(self.num_actions,), + dtype=np.float32, + ) + + @property + def dofs_idx(self) -> list[int]: + """ + Get the indices of the DOF that are enabled (via joint_names). + """ + return self._actuator_manager.dofs_idx + + @property + def pos_dofs_idx(self) -> list[int]: + """ + Get the indices of the DOF that are position controlled. + """ + return self._actuator_manager.pos_dofs_idx + + @property + def vel_dofs_idx(self) -> list[int]: + """ + Get the indices of the DOF that are velocity controlled. + """ + return self._actuator_manager.vel_dofs_idx + + @property + def force_dofs_idx(self) -> list[int]: + """ + Get the indices of the DOF that are force controlled. + """ + return self._actuator_manager.force_dofs_idx + + @property + def default_dofs_pos(self) -> torch.Tensor: + """ + Return the default DOF positions. + """ + return self._actuator_manager.default_dofs_pos + + """ + DOF Getters + """ + + @deprecated( + version="0.3,0", + reason="Use the actuator manager directly.", + ) + def get_dofs_position(self, noise: float = 0.0): + """ + Deprecated: Use the actuator manager directly. + + Return the current position of the enabled DOFs. + This is a wrapper for `RigidEntity.get_dofs_position`. + + Args: + noise: The maximum amount of random noise to add to the position values returned. + """ + return self._actuator_manager.get_dofs_position(noise) + + @deprecated( + version="0.3,0", + reason="Use the actuator manager directly.", + ) + def get_dofs_velocity(self, noise: float = 0.0, clip: tuple[float, float] = None): + """ + Deprecated: Use the actuator manager directly. + + Return the current velocity of the enabled DOFs. + This is a wrapper for `RigidEntity.get_dofs_velocity`. + + Args: + noise: The maximum amount of random noise to add to the velocity values returned. + clip: Clip the velocity returned. + """ + return self._actuator_manager.get_dofs_velocity(noise, clip) + + @deprecated( + version="0.3,0", + reason="Use the actuator manager directly.", + ) + def get_dofs_force(self, noise: float = 0.0, clip_to_max_force: bool = False): + """ + Deprecated: Use the actuator manager directly. + + Return the force experienced by the enabled DOFs. + This is a wrapper for `RigidEntity.get_dofs_force`. + + Args: + noise: The maximum amount of random noise to add to the force values returned. + clip_to_max_force: Clip the force returned to the maximum force defined by the `max_force` parameter. + + Returns: + The force experienced by the enabled DOFs. + """ + return self._actuator_manager.get_dofs_force(noise, clip_to_max_force) + + """ + Operations + """ + + def build(self): + """ + Builds the manager and initialized all the buffers. + """ + + # Define the clip values + lower_limit, upper_limit = self._actuator_manager.get_dofs_limits() + self._clip_values = torch.stack([lower_limit, upper_limit], dim=1) + if self._clip_cfg is not None: + self._get_dof_value_tensor(self._clip_cfg, output=self._clip_values) + + # Scale + self._scale_values = None + if self._scale_cfg is not None: + self._scale_values = self._get_dof_value_tensor(self._scale_cfg) + + # Offset + self._offset_values = None + if self._use_default_offset: + self._offset_values = self._actuator_manager.default_dofs_pos + else: + offset = self._offset_cfg if self._offset_cfg is not None else 0.0 + self._offset_values = self._get_dof_value_tensor(offset) + + def step(self, actions: torch.Tensor) -> torch.Tensor: + """ + Take the incoming actions for this step and handle them. + + Args: + actions: The incoming step actions to handle. + """ + if not self.enabled: + return + actions = super().step(actions) + self._actions = self.handle_actions(actions) + return self._actions + + def handle_actions(self, actions: torch.Tensor) -> torch.Tensor: + """ + Converts the actions to position commands, and send them to the DOF actuators. + Override this function if you want to change the action handling logic. + + Args: + actions: The incoming step actions to handle. + + Returns: + The processed and handled actions. + """ + + # Validate actions + if not self._quiet_action_errors: + if torch.isnan(actions).any(): + print(f"ERROR: NaN actions received! Actions: {actions}") + if torch.isinf(actions).any(): + print(f"ERROR: Infinite actions received! Actions: {actions}") + + # Process pos_actions + pos_dofs=self._actuator_manager.pos_dofs_idx + + pos_actions = actions[pos_dofs] * self._scale_values + self._offset_values + pos_actions = torch.clamp( + pos_actions, + min=self._clip_values[:, 0], + max=self._clip_values[:, 1], + ) + + # Set target positions + self._actuator_manager.control_dofs_position(pos_actions) + + # Process vel_actions + vel_dofs=self._actuator_manager.vel_dofs_idx + + vel_actions = actions[vel_dofs] * self._scale_values + self._offset_values + vel_actions = torch.clamp( + vel_actions, + min=self._clip_values[:, 0], + max=self._clip_values[:, 1], + ) + + # Set target velocities + self._actuator_manager.control_dofs_velocity(vel_actions) + + # Process force_actions + force_dofs=self._actuator_manager.force_dofs_idx + + force_actions = actions[force_dofs] * self._scale_values + self._offset_values + force_actions = torch.clamp( + force_actions, + min=self._clip_values[:, 0], + max=self._clip_values[:, 1], + ) + + # Set target velocities + self._actuator_manager.control_dofs_force(force_actions) + + actions[pos_dofs]=pos_actions + actions[vel_dofs]=vel_actions + actions[force_dofs]=force_actions + return actions + + """ + Internal methods + """ + + def _get_dof_value_tensor( + self, + values: float | dict, + default_value: T = 0.0, + output: torch.Tensor | list[Any] | None = None, + ) -> torch.Tensor: + """ + Given a DofValue dict, loop over the entries, and set the value to the DOF indices (from the actuator) that match the pattern. + + Args: + values: The DOF value to convert (for example: `{".*": 50}`). + + Returns: + A list of values for the DOF indices. + For example, for 4 DOFs: [50, 50, 50, 50] + """ + is_set = [False] * self.num_actions + dof_names = self._actuator_manager.dofs_names + if output is None: + output = torch.zeros( + self.num_actions, device=gs.device, dtype=gs.tc_float + ).fill_(default_value) + for pattern, value in values.items(): + found = False + for i, name in enumerate[str](dof_names): + if not is_set[i] and re.match(f"^{pattern}$", name): + if isinstance(value, (list, tuple)): + value = torch.tensor(value, device=gs.device) + is_set[i] = True + output[i] = value + found = True + if not found: + raise RuntimeError(f"Joint DOF '{pattern}' not found.") + return output diff --git a/genesis_forge/managers/action/position_action_manager.py b/genesis_forge/managers/action/position_action_manager.py index 9fa5a30..9cb5e10 100644 --- a/genesis_forge/managers/action/position_action_manager.py +++ b/genesis_forge/managers/action/position_action_manager.py @@ -325,6 +325,7 @@ def handle_actions(self, actions: torch.Tensor) -> torch.Tensor: # Validate actions if not self._quiet_action_errors: if torch.isnan(actions).any(): + assert NotImplementedError print(f"ERROR: NaN actions received! Actions: {actions}") if torch.isinf(actions).any(): print(f"ERROR: Infinite actions received! Actions: {actions}") diff --git a/genesis_forge/managers/action/velocity_action_manager.py b/genesis_forge/managers/action/velocity_action_manager.py new file mode 100644 index 0000000..0af4075 --- /dev/null +++ b/genesis_forge/managers/action/velocity_action_manager.py @@ -0,0 +1,360 @@ +from __future__ import annotations +import re +import torch +import genesis as gs +import numpy as np +from gymnasium import spaces +from typing import Any, Callable, TypeVar +from deprecated import deprecated +from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.managers.action.base import BaseActionManager +from genesis_forge.values import ensure_dof_pattern +from genesis_forge.managers.actuator import ActuatorManager + +deprecated_arg_names = [ + "joint_names", + "default_pos", + "pd_kp", + "pd_kv", + "max_force", + "damping", + "stiffness", + "frictionloss", + "noise_scale", +] + +T = TypeVar("T") + + +class VelocityActionManager(BaseActionManager): + """ + Converts actions to DOF velocities, using affine transformations (scale). + + .. math:: + + velocity = scaling * action + + Args: + env: The environment to manage the DOF actuators for. + actuator_manager: The actuator manager which is used to setup and control the DOF joints. + scale: How much to scale the action. + clip: Clip the action values to the range. If omitted, the action values will automatically be clipped to the joint limits. + quiet_action_errors: Whether to quiet action errors. + delay_step: The number of steps to delay the actions for. + This is an easy way to emulate the latency in the system. + + Example:: + + class MyEnv(ManagedEnvironment): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # ...define scene and robot... + + def config(self): + self.actuator_manager = ActuatorManager( + self, + joint_names=".*", + default_pos={".*": 0.0}, + kp=50, + kv=0.5, + max_force=8.0, + ) + self.action_manager = VelocityActionManager( + self, + scale=0.5, + actuator_manager=self.actuator_manager, + ) + + Example using the manager directly:: + + class MyEnv(GenesisEnv): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # ...define scene and robot... + + self.actuator_manager = ActuatorManager( + self, + joint_names=".*", + default_pos={".*": 0.0}, + kp=50, + kv=0.5, + max_force=8.0, + ) + self.action_manager = VelocityActionManager( + self, + scale=0.5, + ) + + def build(self): + super().build() + self.actuator_manager.build() + self.action_manager.build() + + step(self, actions: torch.Tensor) -> None: + super().step(actions) + self.action_manager.step(actions) + + # ...do other step things... + + reset(self, envs_idx: list[int] = None) -> None: + super().reset(envs_idx) + self.actuator_manager.reset(envs_idx) + self.action_manager.reset(envs_idx) + + # ...do other reset things... + + + """ + + def __init__( + self, + env: GenesisEnv, + actuator_manager: ActuatorManager | None = None, + scale: float | dict[str, float] = 1.0, + clip: tuple[float, float] | dict[str, tuple[float, float]] = None, + action_handler: Callable[[torch.Tensor], None] = None, + quiet_action_errors: bool = False, + delay_step: int = 0, + **kwargs, + ): + super().__init__(env, delay_step) + self._scale_cfg = ensure_dof_pattern(scale) + self._clip_cfg = ensure_dof_pattern(clip) + self._quiet_action_errors = quiet_action_errors + self._enabled_dof = None + self._actuator_manager = actuator_manager + + self._dofs_pos_buffer: torch.Tensor = None + + # Deprecated actuator parameters + deprecated_actuator_args = { + key: kwargs[key] for key in deprecated_arg_names if key in kwargs + } + if len(deprecated_actuator_args) > 0: + dep_list = ", ".join(deprecated_actuator_args.keys()) + if self._actuator_manager is not None: + raise ValueError( + f"Cannot set both actuator_manager and deprecated actuator parameters: {dep_list}" + ) + print( + f"Actuator arguments are deprecated in the action manager, instead define an ActuatorManager ({dep_list})" + ) + self._actuator_manager = ActuatorManager( + env, + joint_names=kwargs.get("joint_names", ".*"), + default_pos=kwargs.get("default_pos", {".*": 0.0}), + kp=kwargs.get("pd_kp", None), + kv=kwargs.get("pd_kv", None), + max_force=kwargs.get("max_force", None), + damping=kwargs.get("damping", None), + stiffness=kwargs.get("stiffness", None), + frictionloss=kwargs.get("frictionloss", None), + default_noise_scale=kwargs.get("noise_scale", 0.0), + ) + if self._actuator_manager is None: + raise ValueError("No ActuatorManager provided.") + + """ + Properties + """ + + @property + def actuators(self) -> ActuatorManager: + """ + Get the actuator manager. + """ + return self._actuator_manager + + @property + def num_actions(self) -> int: + """ + Get the number of actions. + """ + return self._actuator_manager.num_dofs + + @property + def action_space(self) -> tuple[float, float]: + """ + Returns the actions space for the environment, based on the number of DOFs defined in this action manager. + """ + return spaces.Box( + low=-np.inf, + high=np.inf, + shape=(self.num_actions,), + dtype=np.float32, + ) + + @property + def dofs_idx(self) -> list[int]: + """ + Get the indices of the DOF that are enabled (via joint_names). + """ + return self._actuator_manager.dofs_idx + + @property + def default_dofs_pos(self) -> torch.Tensor: + """ + Return the default DOF positions. + """ + return self._actuator_manager.default_dofs_pos + + """ + DOF Getters + """ + + @deprecated( + version="0.3,0", + reason="Use the actuator manager directly.", + ) + def get_dofs_position(self, noise: float = 0.0): + """ + Deprecated: Use the actuator manager directly. + + Return the current position of the enabled DOFs. + This is a wrapper for `RigidEntity.get_dofs_position`. + + Args: + noise: The maximum amount of random noise to add to the position values returned. + """ + return self._actuator_manager.get_dofs_position(noise) + + @deprecated( + version="0.3,0", + reason="Use the actuator manager directly.", + ) + def get_dofs_velocity(self, noise: float = 0.0, clip: tuple[float, float] = None): + """ + Deprecated: Use the actuator manager directly. + + Return the current velocity of the enabled DOFs. + This is a wrapper for `RigidEntity.get_dofs_velocity`. + + Args: + noise: The maximum amount of random noise to add to the velocity values returned. + clip: Clip the velocity returned. + """ + return self._actuator_manager.get_dofs_velocity(noise, clip) + + @deprecated( + version="0.3,0", + reason="Use the actuator manager directly.", + ) + def get_dofs_force(self, noise: float = 0.0, clip_to_max_force: bool = False): + """ + Deprecated: Use the actuator manager directly. + + Return the force experienced by the enabled DOFs. + This is a wrapper for `RigidEntity.get_dofs_force`. + + Args: + noise: The maximum amount of random noise to add to the force values returned. + clip_to_max_force: Clip the force returned to the maximum force defined by the `max_force` parameter. + + Returns: + The force experienced by the enabled DOFs. + """ + return self._actuator_manager.get_dofs_force(noise, clip_to_max_force) + + """ + Operations + """ + + def build(self): + """ + Builds the manager and initialized all the buffers. + """ + + # Define the clip values + lower_limit, upper_limit = self._actuator_manager.get_dofs_limits() + self._clip_values = torch.stack([lower_limit, upper_limit], dim=1) + if self._clip_cfg is not None: + self._get_dof_value_tensor(self._clip_cfg, output=self._clip_values) + + # Scale + self._scale_values = None + if self._scale_cfg is not None: + self._scale_values = self._get_dof_value_tensor(self._scale_cfg) + + def step(self, actions: torch.Tensor) -> torch.Tensor: + """ + Take the incoming actions for this step and handle them. + + Args: + actions: The incoming step actions to handle. + """ + if not self.enabled: + return + actions = super().step(actions) + self._actions = self.handle_actions(actions) + return self._actions + + def handle_actions(self, actions: torch.Tensor) -> torch.Tensor: + """ + Converts the actions to position commands, and send them to the DOF actuators. + Override this function if you want to change the action handling logic. + + Args: + actions: The incoming step actions to handle. + + Returns: + The processed and handled actions. + """ + + # Validate actions + if not self._quiet_action_errors: + if torch.isnan(actions).any(): + print(f"ERROR: NaN actions received! Actions: {actions}") + if torch.isinf(actions).any(): + print(f"ERROR: Infinite actions received! Actions: {actions}") + + # Process actions + actions = actions * self._scale_values + actions = torch.clamp( + actions, + min=self._clip_values[:, 0], + max=self._clip_values[:, 1], + ) + + # Set target positions + self._actuator_manager.control_dofs_position(actions) + + return actions + + """ + Internal methods + """ + + def _get_dof_value_tensor( + self, + values: float | dict, + default_value: T = 0.0, + output: torch.Tensor | list[Any] | None = None, + ) -> torch.Tensor: + """ + Given a DofValue dict, loop over the entries, and set the value to the DOF indices (from the actuator) that match the pattern. + + Args: + values: The DOF value to convert (for example: `{".*": 50}`). + + Returns: + A list of values for the DOF indices. + For example, for 4 DOFs: [50, 50, 50, 50] + """ + is_set = [False] * self.num_actions + dof_names = self._actuator_manager.dofs_names + if output is None: + output = torch.zeros( + self.num_actions, device=gs.device, dtype=gs.tc_float + ).fill_(default_value) + for pattern, value in values.items(): + found = False + for i, name in enumerate[str](dof_names): + if not is_set[i] and re.match(f"^{pattern}$", name): + if isinstance(value, (list, tuple)): + value = torch.tensor(value, device=gs.device) + is_set[i] = True + output[i] = value + found = True + if not found: + raise RuntimeError(f"Joint DOF '{pattern}' not found.") + return output diff --git a/genesis_forge/managers/actuator/actuator_manager.py b/genesis_forge/managers/actuator/actuator_manager.py index 1daf5f1..adad012 100644 --- a/genesis_forge/managers/actuator/actuator_manager.py +++ b/genesis_forge/managers/actuator/actuator_manager.py @@ -92,6 +92,7 @@ def __init__( env: GenesisEnv, joint_names: list[str] | str = ".*", default_pos: float | NoisyValue | dict = {".*": 0.0}, + control_type: str | NoisyValue | dict = {".*": "position"}, kp: float | NoisyValue | dict = None, kv: float | NoisyValue | dict = None, max_force: float | NoisyValue | tuple[Any, Any] | dict = None, @@ -101,11 +102,13 @@ def __init__( armature: float | NoisyValue | dict = None, default_noise_scale: float = 0.0, entity_attr: str = "robot", + ros_node=None, ): super().__init__(env, type="actuator") self._dofs: dict[str, int] = {} self._robot: RigidEntity = getattr(env, entity_attr) self._default_pos_cfg = ensure_dof_pattern(default_pos) + self._control_type_cfg = ensure_dof_pattern(control_type) self._kp_cfg = ensure_dof_pattern(kp) self._kv_cfg = ensure_dof_pattern(kv) self._max_force_cfg = ensure_dof_pattern(max_force) @@ -114,11 +117,29 @@ def __init__( self._frictionloss_cfg = ensure_dof_pattern(frictionloss) self._armature_cfg = ensure_dof_pattern(armature) self._default_noise_scale = default_noise_scale - - self._batch_dofs_enabled = ( - env.scene.rigid_options.batch_dofs_info - and env.scene.rigid_options.batch_links_info - ) + + if ros_node is not None: + from sensor_msgs.msg import JointState + from builtin_interfaces.msg import Time, Clock + self._ros_node = ros_node + self._ros_clock=Clock() + self._setup_joint_action_publisher() + self._setup_joint_state_subscriber() + else: + self._ros_node =None + + self._pos_actions=[] + self._vel_actions=[] + self._force_actions=[] + self._pos_states=[] + self._vel_states=[] + self._force_states=[] + + if ros_node is None: + self._batch_dofs_enabled = ( + env.scene.rigid_options.batch_dofs_info + and env.scene.rigid_options.batch_links_info + ) self._values: ValueBuffers = { "default_pos": None, @@ -154,6 +175,57 @@ def dofs_idx(self) -> list[int]: Get the indices of the DOF that are enabled (via joint_names). """ return list[int](self._dofs.values()) + + @property + def pos_dofs_idx(self) -> list[int]: + """ + Get the indices of the DOF that are enabled (via joint_names). + """ + pos_dofs_idx = [] + for name, idx in self._dofs.items(): + control_type = None + for pattern, value in self._control_type_cfg.items(): + if re.match(f"^{pattern}$", name): + control_type = value + break + + if control_type == "position": + pos_dofs_idx.append(idx) + return pos_dofs_idx + + @property + def vel_dofs_idx(self) -> list[int]: + """ + Get the indices of the DOF that are enabled (via joint_names). + """ + vel_dofs_idx = [] + for name, idx in self._dofs.items(): + control_type = None + for pattern, value in self._control_type_cfg.items(): + if re.match(f"^{pattern}$", name): + control_type = value + break + + if control_type == "velocity": + vel_dofs_idx.append(idx) + return vel_dofs_idx + + @property + def force_dofs_idx(self) -> list[int]: + """ + Get the indices of the DOF that are enabled (via joint_names). + """ + force_dofs_idx = [] + for name, idx in self._dofs.items(): + control_type = None + for pattern, value in self._control_type_cfg.items(): + if re.match(f"^{pattern}$", name): + control_type = value + break + + if control_type == "force": + force_dofs_idx.append(idx) + return force_dofs_idx @property def dofs_names(self) -> list[str]: @@ -161,6 +233,57 @@ def dofs_names(self) -> list[str]: Get the names of the configured DOFs. """ return list[str](self._dofs.keys()) + + @property + def pos_dofs_names(self) -> list[str]: + """ + Get the names of the configured DOFs with position control. + """ + pos_dofs_names = [] + for name, _ in self._dofs.items(): + control_type = None + for pattern, value in self._control_type_cfg.items(): + if re.match(f"^{pattern}$", name): + control_type = value + break + + if control_type == "position": + pos_dofs_names.append(name) + return pos_dofs_names + + @property + def vel_dofs_names(self) -> list[str]: + """ + Get the names of the configured DOFs with position control. + """ + vel_dofs_names = [] + for name, _ in self._dofs.items(): + control_type = None + for pattern, value in self._control_type_cfg.items(): + if re.match(f"^{pattern}$", name): + control_type = value + break + + if control_type == "velocity": + vel_dofs_names.append(name) + return vel_dofs_names + + @property + def force_dofs_names(self) -> list[str]: + """ + Get the names of the configured DOFs with position control. + """ + force_dofs_names = [] + for name, _ in self._dofs.items(): + control_type = None + for pattern, value in self._control_type_cfg.items(): + if re.match(f"^{pattern}$", name): + control_type = value + break + + if control_type == "velocity": + force_dofs_names.append(name) + return force_dofs_names @property def default_dofs_pos(self) -> torch.Tensor: @@ -170,12 +293,98 @@ def default_dofs_pos(self) -> torch.Tensor: return self._values.get("default_pos", {}).get("buffer", None) @property - def join_names(self) -> list[str]: + def joint_names(self) -> list[str]: """ Get the names of the joints that are enabled, in the order of the DOF indices. """ return list[str](self._dofs.keys()) - + + @property + def propeller_links_idx(self) -> list[int]: + """ + Get the link_idxs of the propeller links + """ + return list(self._robot._propellers_link_idxs) + + @property + def num_propellers(self) -> int: + """ + Get the number of propellers of the robot + """ + return int(self._n_propellers) + + """ + Ros functions and helpers + """ + def _current_sim_timestep(self): + """ + Get the current sim time + """ + if self._ros_node is not None: + timestamp=self._ros_clock.now() + return timestamp + return None + + def _setup_joint_action_publisher(self): + if self._ros_node is not None: + gs.logger.info("Joint actions Publisher started") + def joint_action_callback(js_publisher): + # pos_dof_names,pos_dof_idx_local=self.joint_names,self.dofs_idx + joint_state_msg=JointState() + joint_state_msg.header.stamp=self._current_sim_timestep() + joint_state_msg.name=self.joint_names + joint_state_msg.position=self._pos_actions + joint_state_msg.velocity=self._vel_actions + joint_state_msg.effort=self._force_actions + js_publisher.publish(joint_state_msg) + self.joint_state_publisher = self._ros_node.create_publisher(JointState, f'/joint_commands', 50) + self.timer = self._ros_node.create_timer(0.01, lambda: joint_action_callback(self.joint_state_publisher)) + + def _setup_joint_state_subscriber(self): + if self._ros_node is not None: + gs.logger.info("Joint state Subscriber started") + def joint_state_callback(msg,joint_properties): + motor_dofs=self.dofs_idx + dof_idx_table={} + for k,motor_dof in enumerate(motor_dofs): + dof_idx_table[msg.name[k]]=motor_dof + valid=True + pos_i,vel_i,eff_i=0,0,0 + pos_vals,pos_dofs=[],[] + vel_vals,vel_dofs=[],[] + eff_vals,eff_dofs=[],[] + for joint,joint_contol_type in joint_properties.items(): + if joint_contol_type =="position": + pos_vals.append(msg.position[pos_i]) + pos_dofs.append(dof_idx_table[joint]) + pos_i+=1 + elif joint_contol_type =="velocity": + vel_vals.append(msg.velocity[vel_i]) + vel_dofs.append(dof_idx_table[joint]) + vel_i+=1 + elif joint_contol_type =='force': + eff_vals.append(msg.effort[eff_i]) + eff_dofs.append(dof_idx_table[joint]) + eff_i+=1 + else: + gs.logger.warning("Invalid joint command type") + raise ValueError("Invalid joint command type") + if valid: + self._pos_state=pos_vals + self._vel_state=vel_vals + self._force_state=eff_vals + joint_properties = {} + for name, _ in self._dofs.items(): + control_type = None + for pattern, value in self._control_type_cfg.items(): + if re.match(f"^{pattern}$", name): + control_type = value + break + + joint_properties[name] = control_type + self._ros_node.create_subscription(JointState, + f'joint_states', + lambda: joint_state_callback(joint_properties),100) """ Actuator handlers """ @@ -188,7 +397,10 @@ def get_dofs_position(self, noise: float = 0.0): Args: noise: The maximum amount of random noise to add to the position values returned. """ - pos = self._robot.get_dofs_position(self.dofs_idx) + if self._ros_node is not None and self.env.num_envs==1: + pos=self._pos_states + else: + pos = self._robot.get_dofs_position(self.dofs_idx) if noise > 0.0: pos = self._add_random_noise(pos, noise) return pos @@ -202,7 +414,10 @@ def get_dofs_velocity(self, noise: float = 0.0, clip: tuple[float, float] = None noise: The maximum amount of random noise to add to the velocity values returned. clip: Clip the velocity returned. """ - vel = self._robot.get_dofs_velocity(self.dofs_idx) + if self._ros_node is not None and self.env.num_envs==1: + vel=self._vel_states + else: + vel = self._robot.get_dofs_velocity(self.dofs_idx) if noise > 0.0: vel = self._add_random_noise(vel, noise) if clip is not None: @@ -221,7 +436,10 @@ def get_dofs_force(self, noise: float = 0.0, clip_to_max_force: bool = False): Returns: The force experienced by the enabled DOFs. """ - force = self._robot.get_dofs_force(self.dofs_idx) + if self._ros_node is not None and self.env.num_envs==1: + force=self._force_states + else: + force = self._robot.get_dofs_force(self.dofs_idx) if noise > 0.0: force = self._add_random_noise(force, noise) if clip_to_max_force: @@ -239,6 +457,17 @@ def get_dofs_limits(self) -> tuple[torch.Tensor, torch.Tensor]: Each tensor is of shape (num_envs, num_dofs). """ return self._robot.get_dofs_limit(self.dofs_idx) + + def get_dofs_force_limits(self) -> tuple[torch.Tensor, torch.Tensor]: + """ + Return the force limits of the configured DOFs. + This is a wrapper for `RigidEntity.get_dofs_force_range`. + + Returns: + A tuple of two tensors, the first is the lower limits and the second is the upper limits. + Each tensor is of shape (num_envs, num_dofs). + """ + return self._robot.get_dofs_force_range(self.dofs_idx) def set_dofs_position(self, position: torch.Tensor): """ @@ -260,7 +489,98 @@ def control_dofs_position(self, position: torch.Tensor): position: The position to set the DOFs to. The indices of this tensor should match the configured DOFs (see: `dofs_names` and `dofs_idx` properties). """ - self._robot.control_dofs_position(position, self.dofs_idx) + if self.pos_dofs_idx is not None and len(self.pos_dofs_idx) > 0: + self._robot.control_dofs_position(position, self.pos_dofs_idx) + + def publish_dofs_position(self, position: torch.Tensor): + """ + Publish the position of the configured DOFs in a joint state message. + + Args: + position: The position to set the DOFs to. The indices of this tensor should match the configured DOFs + (see: `dofs_names` and `dofs_idx` properties). + """ + if self.pos_dofs_idx is not None and len(self.pos_dofs_idx) > 0: + self._pos_actions=position.tolist() + + def set_dofs_velocity(self, velocity: torch.Tensor): + """ + Set the velocity of the configured DOFs. + This is a wrapper for `RigidEntity.set_dofs_velocity`. + + Args: + position: The position to set the DOFs to. The indices of this tensor should match the configured DOFs + (see: `dofs_names` and `dofs_idx` properties). + """ + self._robot.set_dofs_velocity(velocity, self.dofs_idx) + + def control_dofs_velocity(self, velocity: torch.Tensor): + """ + Control the velocity of the configured DOFs. + This is a wrapper for `RigidEntity.control_dofs_velocity`. + + Args: + velocity: The velocity to set the DOFs to. The indices of this tensor should match the configured DOFs + (see: `dofs_names` and `dofs_idx` properties). + """ + if self.vel_dofs_idx is not None and len(self.vel_dofs_idx) > 0: + self._robot.control_dofs_velocity(velocity, self.vel_dofs_idx) + + def publish_dofs_velocity(self, velocity: torch.Tensor): + """ + Publish the velocity of the configured DOFs in a joint state message. + + Args: + position: The position to set the DOFs to. The indices of this tensor should match the configured DOFs + (see: `dofs_names` and `dofs_idx` properties). + """ + if self.vel_dofs_idx is not None and len(self.vel_dofs_idx) > 0: + self._vel_actions=velocity.tolist() + + def set_dofs_force(self, force: torch.Tensor): + """ + Set the force of the configured DOFs. + This is a wrapper for `RigidEntity.set_dofs_force`. + + Args: + position: The position to set the DOFs to. The indices of this tensor should match the configured DOFs + (see: `dofs_names` and `dofs_idx` properties). + """ + self._robot.set_dofs_force(force, self.dofs_idx) + + def control_dofs_force(self, force: torch.Tensor): + """ + Control the force of the configured DOFs. + This is a wrapper for `RigidEntity.control_dofs_force`. + + Args: + force: The force to set the DOFs to. The indices of this tensor should match the configured DOFs + (see: `dofs_names` and `dofs_idx` properties). + """ + if self.vel_dofs_idx is not None and len(self.vel_dofs_idx) > 0: + self._robot.control_dofs_force(force, self.vel_dofs_idx) + + def publish_dofs_force(self, force: torch.Tensor): + """ + Publish the force of the configured DOFs in a joint state message. + + Args: + position: The position to set the DOFs to. The indices of this tensor should match the configured DOFs + (see: `dofs_names` and `dofs_idx` properties). + """ + if self.vel_dofs_idx is not None and len(self.vel_dofs_idx) > 0: + self._vel_actions=force.tolist() + + def set_propellels_rpm(self, rpm: torch.Tensor): + """ + Set the propellers's rpm + + Args: + rpm: The rpm to set the propellers to. This function expects the rpm for all the propellers + """ + self._robot.set_propellels_rpm(rpm) + + """ Lifecycle operations @@ -272,7 +592,7 @@ def build(self): """ # Find all configured joints by names/patterns for joint in self._robot.joints: - if joint.type != gs.JOINT_TYPE.REVOLUTE: + if joint.type != gs.JOINT_TYPE.REVOLUTE and joint.type != gs.JOINT_TYPE.PRISMATIC: continue name = joint.name for pattern in self._joint_name_cfg: 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/command_manager.py b/genesis_forge/managers/command/command_manager.py index 433c517..9a01c84 100644 --- a/genesis_forge/managers/command/command_manager.py +++ b/genesis_forge/managers/command/command_manager.py @@ -5,6 +5,7 @@ import genesis as gs from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.ros2_env import Ros2Env from genesis_forge.managers.base import BaseManager from genesis_forge.gamepads import Gamepad @@ -62,7 +63,7 @@ def config(self): def __init__( self, - env: GenesisEnv, + env: GenesisEnv| Ros2Env, range: CommandRange, resample_time_sec: float = 5.0, ): diff --git a/genesis_forge/managers/command/pose_command.py b/genesis_forge/managers/command/pose_command.py new file mode 100644 index 0000000..309eda9 --- /dev/null +++ b/genesis_forge/managers/command/pose_command.py @@ -0,0 +1,246 @@ +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.ros2_env import Ros2Env +from genesis_forge.gamepads import Gamepad + +from .command_manager import CommandManager, CommandRangeValue + + +THIS_DIR = os.path.dirname(os.path.abspath(__file__)) + + +class PoseCommandRange(TypedDict): + pos_x: CommandRangeValue + pos_y: CommandRangeValue + pos_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 commanded velocity arrow""" + + + +DEFAULT_VISUALIZER_CONFIG: PoseDebugVisualizerConfig = { + "envs_idx": None, + "sphere_offset": 0.03, + "sphere_radius": 0.02, + "commanded_color": (0.0, 0.5, 0.0, 1.0), +} + + +class PoseCommandManager(CommandManager): + """ + Generates a position command from uniform distribution. + The command comprises of a linear velocity in x and y direction and an angular velocity around the z-axis. + + 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 + + Arrow 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 arrow visualization + debug_visualizer_cfg: The configuration for the debug visualizer + + Example:: + + class MyEnv(GenesisEnv): + def config(self): + # Create a velocity command manager + self.command_manager = PoseCommandManager( + self, + visualize=True, + range = { + "lin_vel_x_range": (-1.0, 1.0), + "lin_vel_y_range": (-1.0, 1.0), + "ang_vel_z_range": (-0.5, 0.5), + } + ) + + RewardManager( + self, + logging_enabled=True, + cfg={ + "tracking_lin_vel": { + "weight": 1.0, + "fn": rewards.command_tracking_lin_vel, + "params": { + "vel_cmd_manager": self.velocity_command, + }, + }, + "tracking_ang_vel": { + "weight": 1.0, + "fn": rewards.command_tracking_ang_vel, + "params": { + "vel_cmd_manager": self.velocity_command, + }, + }, + # ... other rewards ... + }, + ) + + # Observations + ObservationManager( + self, + cfg={ + "velocity_cmd": {"fn": self.velocity_command.observation}, + # ... other observations ... + }, + ) + """ + + def __init__( + self, + env: GenesisEnv|Ros2Env, + 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._sphere_nodes: list = [] + self.debug_visualizer = debug_visualizer + self.visualizer_cfg = {**DEFAULT_VISUALIZER_CONFIG, **debug_visualizer_cfg} + self.debug_envs_idx = None + + self._is_standing_env = torch.zeros( + env.num_envs, dtype=torch.bool, device=gs.device + ) + + """ + Lifecycle Operations + """ + + def resample_command(self, env_ids: list[int]): + """ + Overwrites commands for environments that should be standing still. + """ + super().resample_command(env_ids) + if not self.enabled: + return + + 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 self.env.scene is not None: + 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() + if self.env.scene is not None: + self._render_arrow() + + def use_gamepad( + self, + gamepad: Gamepad, + pos_x_axis: int = 0, + pos_y_axis: int = 1, + pos_z_axis: int = 2, + euler_x_axis: int = 3, + euler_y_axis: int = 4, + euler_z_axis: int = 5, + ): + """ + Use a connected gamepad to control the command. + + Args: + gamepad: The gamepad to use. + pos_x_axis: Map this gamepad axis index to the position in the x-direction. + pos_y_axis: Map this gamepad axis index to the position in the y-direction. + pos_z_axis: Map this gamepad axis index to the position in the z-direction. + euler_x_axis: Map this gamepad axis index to the euler in the x-direction. + euler_y_axis: Map this gamepad axis index to the euler in the y-direction. + euler_z_axis: Map this gamepad axis index to the euler in the z-direction. + """ + super().use_gamepad( + gamepad, + range_axis={ + "pos_x": pos_x_axis, + "pos_y": pos_y_axis, + "pos_z": pos_z_axis, + "euler_x": euler_x_axis, + "euler_y": euler_y_axis, + "euler_z": euler_z_axis, + }, + ) + + """ + Internal Implementation + """ + + def _render_arrow(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 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..8c2640d --- /dev/null +++ b/genesis_forge/managers/command/position_command.py @@ -0,0 +1,235 @@ +from typing import Tuple, TypedDict + +import os +import torch +import genesis as gs + +from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.ros2_env import Ros2Env +from genesis_forge.gamepads import Gamepad + +from .command_manager import CommandManager, CommandRangeValue + + +THIS_DIR = os.path.dirname(os.path.abspath(__file__)) + + +class PositionCommandRange(TypedDict): + pos_x: CommandRangeValue + pos_y: CommandRangeValue + pos_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 arrows from the top of the robot""" + + sphere_radius: float + """The radius of the shaft of the debug arrows""" + + commanded_color: Tuple[float, float, float, float] + """The color of the commanded velocity arrow""" + + + +DEFAULT_VISUALIZER_CONFIG: PositionDebugVisualizerConfig = { + "envs_idx": None, + "sphere_offset": 0.03, + "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 linear velocity in x and y direction and an angular velocity around the z-axis. + + 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 + + Arrow 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 arrow visualization + debug_visualizer_cfg: The configuration for the debug visualizer + + Example:: + + class MyEnv(GenesisEnv): + def config(self): + # Create a velocity command manager + self.command_manager = PositionCommandManager( + self, + visualize=True, + range = { + "lin_vel_x_range": (-1.0, 1.0), + "lin_vel_y_range": (-1.0, 1.0), + "ang_vel_z_range": (-0.5, 0.5), + } + ) + + RewardManager( + self, + logging_enabled=True, + cfg={ + "tracking_lin_vel": { + "weight": 1.0, + "fn": rewards.command_tracking_lin_vel, + "params": { + "vel_cmd_manager": self.velocity_command, + }, + }, + "tracking_ang_vel": { + "weight": 1.0, + "fn": rewards.command_tracking_ang_vel, + "params": { + "vel_cmd_manager": self.velocity_command, + }, + }, + # ... other rewards ... + }, + ) + + # Observations + ObservationManager( + self, + cfg={ + "velocity_cmd": {"fn": self.velocity_command.observation}, + # ... other observations ... + }, + ) + """ + + def __init__( + self, + env: GenesisEnv|Ros2Env, + 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: list = [] + self.debug_visualizer = debug_visualizer + self.visualizer_cfg = {**DEFAULT_VISUALIZER_CONFIG, **debug_visualizer_cfg} + self.debug_envs_idx = None + + self._is_standing_env = torch.zeros( + env.num_envs, dtype=torch.bool, device=gs.device + ) + + """ + Lifecycle Operations + """ + + def resample_command(self, env_ids: list[int]): + """ + Overwrites commands for environments that should be standing still. + """ + super().resample_command(env_ids) + if not self.enabled: + return + + 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 self.env.scene is not None: + 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() + if self.env.scene is not None: + self._render_sphere() + + def use_gamepad( + self, + gamepad: Gamepad, + pos_x_axis: int = 0, + pos_y_axis: int = 1, + pos_z_axis: int = 2, + ): + """ + Use a connected gamepad to control the command. + + Args: + gamepad: The gamepad to use. + pos_x_axis: Map this gamepad axis index to the position in the x-direction. + pos_y_axis: Map this gamepad axis index to the position in the y-direction. + pos_z_axis: Map this gamepad axis index to the position in the z-direction. + """ + super().use_gamepad( + gamepad, + range_axis={ + "pos_x": pos_x_axis, + "pos_y": pos_y_axis, + "pos_z": pos_z_axis, + }, + ) + + """ + 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 arrows + for sphere in self._sphere_nodes: + self.env.scene.clear_debug_object(sphere) + self._sphere_nodes = [] + + for i in self.debug_envs_idx: + # Target arrow (robot-relative command transformed to world coordinates for visualization) + self._draw_sphere( + pos=self.command[i], + color=self.visualizer_cfg["commanded_color"], + ) + + def _draw_sphere( + self, + pos: torch.Tensor, + color: list[float], + ): + try: + node = self.env.scene.draw_debug_sphere( + pos=pos.cpu().numpy(), + color=color, + radius=self.visualizer_cfg["sphere_radius"], + ) + if node: + self._sphere_nodes.append(node) + 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 53424d1..02ae8ec 100644 --- a/genesis_forge/managers/command/velocity_command.py +++ b/genesis_forge/managers/command/velocity_command.py @@ -5,6 +5,7 @@ import genesis as gs from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.ros2_env import Ros2Env from genesis_forge.utils import entity_lin_vel, transform_by_quat from genesis_forge.gamepads import Gamepad @@ -130,7 +131,7 @@ def config(self): def __init__( self, - env: GenesisEnv, + env: GenesisEnv| Ros2Env, range: VelocityCommandRange, resample_time_sec: float = 5.0, standing_probability: float = 0.0, @@ -172,18 +173,20 @@ def build(self): super().build() # If debug envs_idx is not set, attempt to use the vis_options rendered_envs_idx - 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)) + if self.env.scene is not None: + 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_arrows() + if self.env.scene is not None: + self._render_arrows() def use_gamepad( self, diff --git a/genesis_forge/managers/entity_manager.py b/genesis_forge/managers/entity_manager.py index 2ed68cf..e204abb 100644 --- a/genesis_forge/managers/entity_manager.py +++ b/genesis_forge/managers/entity_manager.py @@ -76,6 +76,7 @@ def __init__( env: GenesisEnv, entity_attr: str, on_reset: dict[str, EntityResetConfig], + ros_node=None ): super().__init__(env, type="entity") if hasattr(env, "add_entity_manager"): @@ -84,6 +85,15 @@ def __init__( self.entity: RigidEntity | None = None self.on_reset = on_reset self._entity_attr = entity_attr + + if ros_node is not None: + from geometry_msgs.msg import Pose + from geometry_msgs.msg import Twist + self._ros_node=ros_node + self._setup_pose_subscriber() + self._setup_twist_subscriber() + else: + self._ros_node=None # Wrap config items self.on_reset: dict[str, ConfigItem] = {} @@ -131,6 +141,20 @@ def inv_base_quat(self) -> torch.Tensor: Helpers """ + def _setup_pose_subscriber(self): + if self._ros_node is not None: + def pose_subscriber_callback(msg): + self._robot_pos=torch.tensor(msg.poisition) + self._robot_quat=torch.tensor(self._to_gs_quat_format(msg.orientation)) + self._ros_node.create_subscription(Pose,"pose_robotic",pose_subscriber_callback,100) + + def _setup_twist_subscriber(self): + if self._ros_node is not None: + def twist_subscriber_callback(msg): + self._robot_lin_vel=torch.tensor(msg.linear) + self._robot_ang_vel=torch.tensor(msg.angular) + self._ros_node.create_subscription(Twist,"twist_robotic",twist_subscriber_callback,100) + def get_projected_gravity(self) -> torch.Tensor: """ The projected gravity of the entity's base link, in the entity's local frame. @@ -141,13 +165,19 @@ def get_linear_velocity(self) -> torch.Tensor: """ The linear velocity of the entity's base link, in the entity's local frame. """ - return transform_by_quat(self.entity.get_vel(), self._inv_base_quat) - + if self._ros_node is not None and self.env.num_envs==1: + return transform_by_quat(self._robot_lin_vel, self._inv_base_quat) + else: + return transform_by_quat(self.entity.get_vel(), self._inv_base_quat) + def get_angular_velocity(self) -> torch.Tensor: """ The angular velocity of the entity's base link, in the entity's local frame. """ - return transform_by_quat(self.entity.get_ang(), self._inv_base_quat) + if self._ros_node is not None and self.env.num_envs==1: + return transform_by_quat(self._robot_ang_vel, self._inv_base_quat) + else: + return transform_by_quat(self.entity.get_ang(), self._inv_base_quat) """ Operations. @@ -194,6 +224,11 @@ def _cached_calcs(self): """ Calculate and cache some common values """ - self._base_pos[:] = self.entity.get_pos() - self._base_quat[:] = self.entity.get_quat() - self._inv_base_quat = inv_quat(self._base_quat) + if self._ros_node is None: + self._base_pos[:] = self.entity.get_pos() + self._base_quat[:] = self.entity.get_quat() + self._inv_base_quat = inv_quat(self._base_quat) + else: + self._base_pos[:] = self._robot_pos + self._base_quat[:] = self._robot_quat + self._inv_base_quat = inv_quat(self._base_quat) diff --git a/genesis_forge/managers/observation_manager.py b/genesis_forge/managers/observation_manager.py index 4826ba1..d2848f0 100644 --- a/genesis_forge/managers/observation_manager.py +++ b/genesis_forge/managers/observation_manager.py @@ -1,5 +1,6 @@ import torch import numpy as np +import collections from gymnasium import spaces import genesis as gs from typing import TypedDict, Callable, Any @@ -136,12 +137,14 @@ def __init__( cfg: dict[str, ObservationConfig], name: str = "policy", history_len: int | None = None, + obs_to_discard: list[str]=None, noise: tuple[float, float] | None = None, ): super().__init__(env, "observation") self._name = name self.cfg = cfg self.noise = noise + self.obs_to_discard=obs_to_discard if obs_to_discard is not None else [] self._observation_size = 1 self._observation_space = None @@ -197,39 +200,52 @@ def build(self): assert callable(cfg.fn), f"Observation function {name} is not callable" # Make an initial observation and create the observation space - obs = self._perform_observation() - single_obs_size = obs.shape[1] - self._observation_size = single_obs_size * self._history_len - self._observation_space = spaces.Box( - low=-np.inf, - high=np.inf, - shape=(self._observation_size,), - dtype=np.float32, - ) - - # Fill history buffer - shape = (self.env.num_envs, single_obs_size) - self._history = [ - torch.zeros(shape, device=gs.device) for _ in range(self._history_len) - ] + obs = self._perform_observation(single_agent=True) + single_obs_size=0 + if len(self.obs_to_discard)>0: + for key,value in obs.items(): + if key not in self.obs_to_discard: + single_obs_size +=value.shape[0] + else: + for key,value in obs.items(): + if value.ndim<=2: + single_obs_size +=value.shape[0] + + self._observation_space=self._spec_to_space(obs) + + if self._history_len>1: + self._observation_size = single_obs_size * self._history_len + + # Fill history buffer + shape = (self.env.num_envs, single_obs_size) + self._history = [ + torch.zeros(shape, device=gs.device) for _ in range(self._history_len) + ] def get_observations(self) -> torch.Tensor: """Generate current observations for all environments.""" if not self.enabled: return torch.zeros((self.env.num_envs, self._observation_size)) - - self._history.pop() + obs = self._perform_observation() - self._history.insert(0, obs) - return torch.cat(self._history, dim=-1) + obs_flattened=torch.cat([ + obs_value.view(self.env.num_envs, -1) if obs_value.dim() != 2 else obs_value + for obs_value in obs.values() + ],dim=-1) + # only use the _history if _history_len is grater than 1 + if self._history_len>1: + self._history.pop() + self._history.insert(0, obs_flattened) + return torch.cat(self._history, dim=-1) + return obs_flattened """ Private methods. """ - def _perform_observation(self) -> torch.Tensor: + def _perform_observation(self,single_agent=False) -> torch.Tensor: """Perform a round of observations.""" - obs = [] + obs = {} for name, cfg in self.cfg.items(): try: # Get values @@ -246,9 +262,32 @@ def _perform_observation(self) -> torch.Tensor: if noise is not None and noise != 0.0: noise_value = torch.empty_like(value).uniform_(-1, 1) * noise value += noise_value - - obs.append(value) + obs[name]=value[0] if single_agent else value except Exception as e: print(f"Error generating observation for '{name}'") raise e - return torch.cat(obs, dim=-1) + return collections.OrderedDict(obs) + + def _spec_to_space(self, spec: Any) -> spaces.Space: + """returns an gymnasium.space action space based on the observation """ + if type(spec) is tuple: + return spaces.Box(shape=spec[0].shape, dtype=np.float64, low=spec[0], high=spec[1]) + elif isinstance(spec, np.ndarray): + return spaces.Box( + shape=spec.shape, + dtype=np.float64, + low=np.full(spec.shape, float("-inf")), + high=np.full(spec.shape, float("inf")), + ) + elif isinstance(spec, torch.Tensor): + spec.detach().cpu().numpy() + return spaces.Box( + shape=spec.shape, + dtype=np.float64, + low=np.full(spec.shape, float("-inf")), + high=np.full(spec.shape, float("inf")), + ) + elif isinstance(spec, collections.OrderedDict): + return spaces.Dict({k: self._spec_to_space(v) for k, v in spec.items()}) + else: + raise ValueError(f"Spec type {type(spec)} not supported. Please report this issue") diff --git a/genesis_forge/managers/sensors/__init__.py b/genesis_forge/managers/sensors/__init__.py new file mode 100644 index 0000000..fa4499d --- /dev/null +++ b/genesis_forge/managers/sensors/__init__.py @@ -0,0 +1,15 @@ +from .imu_manager import ImuManager +from .contact.contact_manager import ContactManager +from .spherical_raycaster_manager import SphericalRaycasterManager +from .grid_raycaster_manager import GridRaycasterManager +from .depth_camera_manager import DepthCameraManager +from .camera_manager import CameraManager + +__all__ = [ + "ImuManager", + "ContactManager", + "GridRaycasterManager", + "SphericalRaycasterManager", + "DepthCameraManager", + "CameraManager" +] diff --git a/genesis_forge/managers/sensors/base_sensor_manager.py b/genesis_forge/managers/sensors/base_sensor_manager.py new file mode 100644 index 0000000..428d9c0 --- /dev/null +++ b/genesis_forge/managers/sensors/base_sensor_manager.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +import torch +import genesis as gs + +from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.managers.base import BaseManager +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from genesis.engine.entities import RigidEntity + +class BaseSensorManager(BaseManager): + """ + Base sensor for handling a sensor. + + Args: + sensor_name: The name of the sensor + env: The environment to sense. + link_name: The name of the link the sensor is attached to + entity_attr: The environment attribute which contains the entity with the links we're tracking. Defaults to `robot`. + pos_offset: The xyz position offset between the link origin and the sensor origin + euler_offset: The xyz orintation offset between the link origin and the sensor origin + delay: The daly induced between sensor measurements + draw_debug: Boolean value indicating whether the sensors data should be visualised in the scene + read_frequency: the frequncy at which to read the input from the sensor in hz + + # ...other arguments here... + """ + #TODO: write proper example + + def __init__( + self, + env: GenesisEnv, + sensor_type: str, + sensor_name: str, + link_name: str, + entity_attr: RigidEntity = "robot", + pos_offset: list[float]=[0,0,0], + euler_offset: list[float]=[0,0,0], + delay: float=0.0, + draw_debug: bool=False, + read_freuency: float=20, + ): + super().__init__(env, type=sensor_type) + self._sensor_name = sensor_name + self._link_name = link_name + self._entity_attr = entity_attr + self._pos_offset=pos_offset + self._euler_offset=euler_offset + self._delay=delay + self._draw_debug=draw_debug + self._read_frequncy=read_freuency + self._sensor_read_interval=1/self._read_frequncy + + # Get the link indices + if self._link_name is None: + self._link_name="base_link" + self._entity,self._link = self._get_entity_and_link( + self._entity_attr, self._link_name + ) + + self.base_sensor_args=dict( + entity_idx=self._entity.idx, + link_idx_loacl=self._link.idx_local, + pos_offset=self._pos_offset, + euler_offset=self._euler_offset + ) + self._last_reading_timestamp=None + + """ + Properties + """ + @property + def sensor_name(self) -> torch.Tensor: + """name of the sensor.""" + return self._sensor_name + + @property + def link_name(self) -> torch.Tensor: + """name of the sensor link.""" + return self._link_name + + @property + def link_idx(self) -> torch.Tensor: + """The link index for the sensor link.""" + return self._link.idx + + @property + def local_link_idx(self) -> torch.Tensor: + """The local link index for the sensor link.""" + return self._link.idx_local + + """ + Lifecycle Operations + """ + + def build(self): + """Initialize link indices and buffers.""" + super().build() + + def reset(self, envs_idx: list[int] | None = None): + super().reset(envs_idx) + if not self.enabled: + return + if envs_idx is None: + self.envs_idx = torch.arange(self.env.num_envs, device=gs.device) + else: + self.envs_idx=envs_idx + + + def step(self): + super().step() + if not self.enabled: + return + + def disable_sensor(self): + self.enabled=False + + """ + Internal Implementation + """ + + def _get_entity_and_link( + self, entity_attr: str, link_name: str + ): + """ + Find the link handle for the given link name and entity_name. + + Args: + entity: The entity to find the links in. + link: The link found in the entity + + Returns: Tuple of global and local link index tensors. + """ + entity = self.env.__getattribute__(entity_attr) + try: + link=entity.get_link(link_name) + except Exception as e: + link=entity.links[0] + + return entity,link + + def __repr__(self): + attrs = [f"sensor_name={self._sensor_name}",f"link_name={self._link_name}"] + if self._entity_attr: + attrs.append(f"entity_attr={self._entity_attr}") + if self._pos_offset: + attrs.append(f"pos_offset={self._pos_offset}") + if self._euler_offset: + attrs.append(f"pos_offset={self._euler_offset}") + if self._delay: + attrs.append(f"delay={self._delay}") + if self._draw_debug: + attrs.append(f"draw_debug={self._draw_debug}") + attrs_str = ", ".join(attrs) + return f"{self.__class__.__name__}({attrs_str})" diff --git a/genesis_forge/managers/sensors/camera_manager.py b/genesis_forge/managers/sensors/camera_manager.py new file mode 100644 index 0000000..73a39ea --- /dev/null +++ b/genesis_forge/managers/sensors/camera_manager.py @@ -0,0 +1,317 @@ +from __future__ import annotations + +import numpy as np +import torch +import genesis as gs +from genesis.utils.geom import quat_to_R, euler_to_quat + +from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.managers.sensors.base_sensor_manager import BaseSensorManager +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from genesis.engine.entities import RigidEntity + +class CameraManager(BaseSensorManager): + """ + A grid Lidar attached to an entity's link in the environment. + + Args: + sensor_name: The name of the sensor + env: The environment to sense. + link_name: The name of the link the sensor is attached to + entity_attr: The environment attribute which contains the entity with the links we're tracking. Defaults to `robot`. + pos_offset: The xyz position offset between the link origin and the sensor origin + euler_offset: The xyz orintation offset between the link origin and the sensor origin + delay: The daly induced between sensor measurements + draw_debug: Boolean value indicating whether the sensors data should be visualised in the scene + read_frequency: the frequncy at which to read the input from the sensor in hz + model : str + Specifies the camera model. Options are 'pinhole' or 'thinlens'. + res : tuple of int, shape (2,) + The resolution of the camera, specified as a tuple (width, height). + fov : float + The vertical field of view of the camera in degrees. + aperture : float + The aperture size of the camera, controlling depth of field. + focus_dist : float | None + The focus distance of the camera. If None, it will be auto-computed using `pos` and `lookat`. + GUI : bool + Whether to display the camera's rendered image in a separate GUI window. + spp : int, optional + Samples per pixel. Only available when using RayTracer renderer. Defaults to 256. + denoise : bool + Whether to denoise the camera's rendered image. Only available when using the RayTracer renderer. Defaults + to True on Linux, otherwise False. If OptiX denoiser is not available in your platform, consider enabling + the OIDN denoiser option when building the RayTracer. + near : float + Distance from camera center to near plane in meters. + Only available when using rasterizer in Rasterizer and BatchRender renderer. Defaults to 0.1. + far : float + Distance from camera center to far plane in meters. + Only available when using rasterizer in Rasterizer and BatchRender renderer. Defaults to 20.0. + env_idx : int, optional + The specific environment index to bind to the camera. This option must be specified if and only if a + non-batched renderer is being used. If provided, only this environment will be taken into account when + following a rigid entity via 'follow_entity' and when being attached to some rigid link via 'attach'. Note + that this option is unrelated to which environment is being rendering on the scene. Default to None for + batched renderers (ie BatchRender), 'rendered_envs_idx[0]' otherwise (ie Raytracer or Rasterizer). + debug_camera : bool + Whether to use the debug camera. It enables to create cameras that can used to monitor / debug the + simulation without being part of the "sensors". Their output is rendered by the usual simple Rasterizer + systematically, no matter if BatchRender and RayTracer is enabled. This way, it is possible to record the + simulation with arbitrary resolution and camera pose, without interfering with what robots can perceive + from their environment. Defaults to False. + render_on_demand: bool + where to only render the camera on the 'get_(camera_type)_img' function call + rgb: bool + whether to render RGB images + depth: bool + whether to render depth images + segmentation: bool + whether to render segmentation images + normal: bool + whether to render Normal images + """ + #TODO: write proper example + + def __init__( + self, + env: GenesisEnv, + sensor_name, + link_name: str, + entity_attr: RigidEntity = "robot", + pos_offset: list[float] = [0,0,0], + euler_offset: list[float] = [0,0,0], + delay: float = 0.0, + draw_debug: bool = False, + read_frequency: float= 15, + model: str = 'pinhole', + res: tuple[int, int] = (640, 480), + fov: float = 45.0, + aperture: float = 0.0, + focus_dist: float | None = None, + spp: int = 256, + denoise: bool = True, + near: float = 0.1, + far: float = 20.0, + env_idx: int | None = None, + debug_camera: bool = False, + render_on_demand=False, + rgb: bool =True, + depth: bool =False, + segmentation: bool=False, + normal: bool=False + + ): + super().__init__( + env=env, + sensor_type="camera_sensor", + sensor_name=sensor_name, + link_name=link_name, + entity_attr=entity_attr, + pos_offset=pos_offset, + euler_offset=euler_offset, + delay=delay, + draw_debug=draw_debug, + read_freuency=read_frequency) + + self._render_on_demand=render_on_demand + self._render_rgb=rgb + self._render_depth=depth + self._render_segmentation=segmentation + self._render_normal=normal + + self._camera_resolution=res + self._camera_sensors=[] + for env_idx in range(self.env.num_envs): + self._camera_sensors.append(self.env.scene.add_camera( + model=model, + res=res, + fov=fov, + aperture=aperture, + focus_dist=focus_dist, + GUI=self._draw_debug, + spp=spp, + denoise=denoise, + near=near, + far=far, + env_idx=env_idx, + debug=debug_camera, + ) + ) + + """ + Properties + """ + + @property + def link_name(self) -> torch.Tensor: + """name of the imu link.""" + return self._link_name + + @property + def link_idx(self) -> torch.Tensor: + """The link index for the imu link.""" + return self._link.idx + + @property + def local_link_idx(self) -> torch.Tensor: + """The local link index for the imu link.""" + return self._link.idx_local + + """ + Helper Methods + """ + def get_rgb_img(self) -> torch.Tensor: + """ + RGB image rendered by the camera sensor + + Returns: + The RGB image shape is (n_envs,cam_res_width,cam_res_height, 3) + """ + if not self._render_on_demand and self._render_rgb: + return self._sensor_reading_rgb + else: + return self._render_all[0] + + def get_depth_img(self) -> torch.Tensor: + """ + depth image rendered by the camera sensor + + Returns: + The depth image shape is (n_envs,cam_res_width,cam_res_height, 1) + """ + if not self._render_on_demand and self._render_depth: + return self._sensor_reading_depth + else: + return self._render_all()[1] + + def get_sgementation_img(self) -> torch.Tensor: + """ + segmentation image rendered by the camera sensor + + Returns: + The segmentation image shape is (n_envs,cam_res_width,cam_res_height, 1) + """ + if not self._render_on_demand and self._render_segmentation: + return self._sensor_reading_segmentation + else: + return self._render_all()[2] + + def get_normal_img(self) -> torch.Tensor: + """ + normal image rendered by the camera sensor + + Returns: + The normal image shape is (n_envs,cam_res_width,cam_res_height, 3) + """ + if not self._render_on_demand and self._render_normal: + return self._sensor_reading_normal + else: + return self._render_all()[3] + + """ + Lifecycle Operations + """ + + def build(self): + """Initialize link indices and buffers.""" + super().build() + T=np.eye(4) + T[:3,:3]=quat_to_R(euler_to_quat(np.array(self._euler_offset))) + T[:3,3]=self._pos_offset + for camera_sensor in self._camera_sensors: + camera_sensor.attach(self._link,offset_T=T) + + if not self._render_on_demand: + if self._render_rgb: + self._sensor_reading_rgb=torch.zeros(self.env.num_envs, + self._camera_resolution[0], + self._camera_resolution[1],3) + if self._render_depth: + self._sensor_reading_depth=torch.zeros(self.env.num_envs, + self._camera_resolution[0], + self._camera_resolution[1],1) + if self._render_segmentation: + self._sensor_reading_segmentation=torch.zeros(self.env.num_envs, + self._camera_resolution[0], + self._camera_resolution[1],1) + if self._render_normal: + self._sensor_reading_normal=torch.zeros(self.env.num_envs, + self._camera_resolution[0], + self._camera_resolution[1],3) + + def reset(self, envs_idx: list[int] | None = None): + super().reset(envs_idx) + if not self.enabled: + return + for camera_sensor in self._camera_sensors: + camera_sensor.move_to_attach() + if not self._render_on_demand: + if self._render_rgb: + self._sensor_reading_rgb[self.envs_idx]=0.0 + if self._render_depth: + self._sensor_reading_depth[self.envs_idx]=0.0 + if self._render_segmentation: + self._sensor_reading_segmentation[self.envs_idx]=0.0 + if self._render_normal: + self._sensor_reading_normal[self.envs_idx]=0.0 + + def step(self): + super().step() + if not self.enabled: + return + if self._last_reading_timestamp is None: + self._last_reading_timestamp=self.env.scene.cur_t + elif self.env.scene.cur_t-self._last_reading_timestamp>self._sensor_read_interval: + self._last_reading_timestamp=self.env.scene.cur_t + if not self._render_on_demand: + rgb,depth,segmentation,normal=self._render_all() + if self._render_rgb: + self._sensor_reading_rgb=rgb + if self._render_depth: + self._sensor_reading_depth=depth + if self._render_segmentation: + self._sensor_reading_segmentation=segmentation + if self._render_normal: + self._sensor_reading_normal=normal + + def _render_all(self): + rgb_all=None + depth_all=None + segmentation_all=None + normal_all=None + for camera_sensor in self._camera_sensors: + rgb,depth,segmentation,normal=camera_sensor.render(rgb=self._render_rgb, + depth=self._render_depth, + segmentation=self._render_segmentation, + normal=self._render_normal) + if self._render_rgb: + if rgb_all is None: + rgb_all=rgb + else: + rgb_all=np.concat([rgb_all,rgb]) + if self._render_depth: + if depth_all is None: + depth_all=depth + else: + depth_all=np.concat([depth_all,depth]) + if self._render_segmentation: + if segmentation_all is None: + segmentation_all=segmentation + else: + segmentation_all=np.concat([segmentation_all,segmentation]) + if self._render_normal: + if normal_all is None: + normal_all=normal + else: + normal_all=np.concat([normal_all,normal]) + return ( + torch.tensor(rgb_all) if self._render_rgb else None, + torch.tensor(depth_all) if self._render_depth else None, + torch.tensor(segmentation_all) if self._render_segmentation else None, + torch.tensor(normal_all) if self._render_normal else None + ) + + diff --git a/genesis_forge/managers/sensors/contact/config.py b/genesis_forge/managers/sensors/contact/config.py new file mode 100644 index 0000000..b3156c2 --- /dev/null +++ b/genesis_forge/managers/sensors/contact/config.py @@ -0,0 +1,25 @@ +from typing import TypedDict, Tuple + + +class ContactDebugVisualizerConfig(TypedDict): + """Defines the configuration for the contact debug visualizer.""" + + envs_idx: list[int] + """The indices of the environments to visualize. If None, all environments will be visualized.""" + + color: Tuple[float, float, float, float] + """The color of the contact ball""" + + radius: float + """The radius of the visualization sphere""" + + force_threshold: float + """The threshold, in Newtons, for the contact force to be visualized""" + + +DEFAULT_VISUALIZER_CONFIG: ContactDebugVisualizerConfig = { + "envs_idx": None, + "size": 0.03, + "color": (0.5, 0.0, 0.0, 1.0), + "force_threshold": 1.0, +} diff --git a/genesis_forge/managers/sensors/contact/contact_manager.py b/genesis_forge/managers/sensors/contact/contact_manager.py new file mode 100644 index 0000000..70132f9 --- /dev/null +++ b/genesis_forge/managers/sensors/contact/contact_manager.py @@ -0,0 +1,551 @@ +from __future__ import annotations + +import re +import torch +import genesis as gs + +from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.managers.base import BaseManager +from genesis_forge.managers.sensors.contact.config import ( + ContactDebugVisualizerConfig, + DEFAULT_VISUALIZER_CONFIG, +) +from genesis_forge.managers.sensors.contact.kernel import kernel_get_contact_forces + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from genesis.engine.entities import RigidEntity + + +class ContactManager(BaseManager): + """ + Tracks the contact forces between entity links in the environment. + + Args: + env: The environment to track the contact forces for. + link_names: The names, or name regex patterns, of the entity links to track the contact forces for. + entity_attr: The environment attribute which contains the entity with the links we're tracking. Defaults to `robot`. + with_entity_attr: Filter the contact forces to only include contacts with the entity assigned to this environment attribute. + with_links_names: Filter the contact forces to only include contacts with these links. + track_air_time: Whether to track the air time of the entity link contacts. + air_time_contact_threshold: When track_air_time is True, this is the threshold for the contact forces to be considered. + debug_visualizer: Whether to visualize the contact points. + debug_visualizer_cfg: The configuration for the contact debug visualizer. + + Example with ManagedEnvironment:: + + class MyEnv(ManagedEnvironment): + + # ... Construct scene and other env setup ... + + def config(self): + # Define contact manager + self.foot_contact_manager = ContactManager( + self, + link_names=[".*_Foot"], + ) + + # Use contact manager in rewards + self.reward_manager = RewardManager( + self, + term_cfg={ + "Foot contact": { + "weight": 5.0, + "fn": rewards.has_contact, + "params": { + "contact_manager": self.foot_contact_manager, + "min_contacts": 4, + }, + }, + }, + ) + + # ... other managers here ... + + Example using the contact manager directly:: + + class MyEnv(GenesisEnv): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.contact_manager = ContactManager( + self, + link_names=[".*_Foot"], + ) + + def build(self): + super().build() + self.contact_manager.build() + + def step(self, actions: torch.Tensor): + super().step(actions) + self.contact_manager.step() + return obs, rewards, terminations, timeouts, info + + def reset(self, envs_idx: list[int] | None = None): + super().reset(envs_idx) + self.contact_manager.reset(envs_idx) + return obs, info + + def calculate_rewards(): + # Reward for each foot in contact with something with at least 1.0N force + CONTACT_THRESHOLD = 1.0 + CONTACT_WEIGHT = 0.005 + has_contact = self.contact_manager.contacts[:,:].norm(dim=-1) > CONTACT_THRESHOLD + contact_reward = has_contact.sum(dim=1).float() * CONTACT_WEIGHT + + # Access contact positions for debugging or additional analysis + contact_positions = self.contact_manager.contact_positions + # contact_positions shape: (n_envs, n_target_links, 3) + # Positions are automatically averaged when multiple contacts occur + + # ...additional reward calculations here... + + Filtering:: + + class MyEnv(ManagedEnvironment): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.scene = gs.Scene() + + # Add terrain + self.terrain = self.scene.add_entity(gs.morphs.Plane()) + + # add robot + self.robot = self.scene.add_entity( + gs.morphs.URDF(file="urdf/go2/urdf/go2.urdf"), + ) + + def config(self): + # Track all contacts between the robot's feet and the terrain + self.contact_manager = ContactManager( + self, + entity_attr="robot", + link_names=[".*_foot"], + with_entity_attr="terrain", + ) + + # ...other managers here... + + # ...other operations here... + """ + + def __init__( + self, + env: GenesisEnv, + link_names: list[str], + entity_attr: RigidEntity = "robot", + with_entity_attr: RigidEntity = None, + with_links_names: list[int] = None, + track_air_time: bool = False, + air_time_contact_threshold: float = 1.0, + debug_visualizer: bool = False, + debug_visualizer_cfg: ContactDebugVisualizerConfig = DEFAULT_VISUALIZER_CONFIG, + ): + super().__init__(env, "contact") + + self._link_names = link_names + self._air_time_contact_threshold = air_time_contact_threshold + self._track_air_time = track_air_time + self._entity_attr = entity_attr + self._link_ids = None + self._local_link_ids = None + self._with_entity_attr = with_entity_attr + self._with_links_names = with_links_names + self._with_link_ids = torch.empty(0, device=gs.device) + self._with_local_link_ids = None + self._has_with_filter = ( + with_entity_attr is not None or with_links_names is not None + ) + + self.debug_visualizer = debug_visualizer + self.debug_envs_idx = None + self.visualizer_cfg = {**DEFAULT_VISUALIZER_CONFIG, **debug_visualizer_cfg} + self._debug_nodes = [] + self._contact_position_counts = None + + self.contacts: torch.Tensor | None = None + """Contact forces experienced by the entity links.""" + + self.contact_positions: torch.Tensor | None = None + """Contact positions for each target link.""" + + self.last_air_time: torch.Tensor | None = None + """Time spent (in s) in the air before the last contact.""" + + self.current_air_time: torch.Tensor | None = None + """Time spent (in s) in the air since the last detach.""" + + self.last_contact_time: torch.Tensor | None = None + """Time spent (in s) in contact before the last detach.""" + + self.current_contact_time: torch.Tensor | None = None + """Time spent (in s) in contact since the last contact.""" + + """ + Properties + """ + + @property + def link_ids(self) -> torch.Tensor: + """The global link indices for the target links.""" + return self._link_ids + + @property + def local_link_ids(self) -> torch.Tensor: + """The local link indices for the target links.""" + return self._local_link_ids + + """ + Helper Methods + """ + + def has_made_contact(self, dt: float, time_margin: float = 1.0e-8) -> torch.Tensor: + """ + Checks if links that have established contact within the last :attr:`dt` seconds. + + This function checks if the links have established contact within the last :attr:`dt` seconds + by comparing the current contact time with the given time period. If the contact time is less + than the given time period, then the links are considered to be in contact. + + Args: + dt: The time period since the contact was established. + time_margin: Adds a little error margin to the dt time period. + + Returns: + A boolean tensor indicating the links that have established contact within the last + :attr:`dt` seconds. Shape is (n_envs, n_target_links) + + Raises: + RuntimeError: If the manager is not configured to track air time. + """ + # check if the sensor is configured to track contact time + if not self._track_air_time: + raise RuntimeError( + "The contact sensor is not configured to track air time." + "Please enable the 'track_air_time' in the manager configuration." + ) + # check if the bodies are in contact + currently_in_contact = self.current_contact_time > 0.0 + less_than_dt_in_contact = self.current_contact_time < (dt + time_margin) + return currently_in_contact * less_than_dt_in_contact + + def has_broken_contact( + self, dt: float, time_margin: float = 1.0e-8 + ) -> torch.Tensor: + """Checks links that have broken contact within the last :attr:`dt` seconds. + + This function checks if the links have broken contact within the last :attr:`dt` seconds + by comparing the current air time with the given time period. If the air time is less + than the given time period, then the links are considered to not be in contact. + + Args: + dt: The time period since the contact was broken. + time_margin: Adds a little error margin to the dt time period. + + Returns: + A boolean tensor indicating the links that have broken contact within the last + :attr:`dt` seconds. Shape is (n_envs, n_target_links) + + Raises: + RuntimeError: If the manager is not configured to track air time. + """ + # check if the sensor is configured to track contact time + if not self._track_air_time: + raise RuntimeError( + "The contact manager is not configured to track air time." + "Please enable the 'track_air_time' in the manager configuration." + ) + currently_detached = self.current_air_time > 0.0 + less_than_dt_detached = self.current_air_time < (dt + time_margin) + return currently_detached * less_than_dt_detached + + def get_contact_forces(self, link_idx: int | list[int]) -> torch.Tensor: + """ + Get the contact forces for one or more links + + Args: + link_idx: The link index or list of link indices to get the contact forces for. + + Returns: + The contact forces for the target links. Shape is (n_envs, n_target_links, 3) + """ + idx = [] + if isinstance(link_idx, int): + idx = torch.nonzero(self._link_ids == link_idx)[0] + elif isinstance(link_idx, list): + idx = [torch.nonzero(self._link_ids == i)[0].item() for i in link_idx] + return self.contacts[:, idx, :] + + """ + Lifecycle Operations + """ + + def build(self): + """Initialize link indices and buffers.""" + super().build() + + # If debug envs_idx is not set, attempt to use the vis_options rendered_envs_idx + 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)) + + # Get the link indices + (self._link_ids, self._local_link_ids) = self._get_links_idx( + self._entity_attr, self._link_names + ) + if not self._link_ids.is_contiguous(): + self._link_ids = self._link_ids.contiguous() + if self._with_entity_attr or self._with_links_names: + with_entity_attr = ( + self._with_entity_attr + if self._with_entity_attr is not None + else "robot" + ) + (self._with_link_ids, self._with_local_link_ids) = self._get_links_idx( + with_entity_attr, self._with_links_names + ) + if not self._with_link_ids.is_contiguous(): + self._with_link_ids = self._with_link_ids.contiguous() + + # Initialize buffers + link_count = self._link_ids.shape[0] + self.contacts = torch.zeros( + (self.env.num_envs, link_count, 3), device=gs.device + ) + self.contact_positions = torch.zeros( + (self.env.num_envs, link_count, 3), device=gs.device + ) + self._contact_position_counts = torch.zeros( + (self.env.num_envs, link_count), device=gs.device + ) + if self._track_air_time: + self.last_air_time = torch.zeros( + (self.env.num_envs, link_count), device=gs.device + ) + self.current_air_time = torch.zeros_like(self.last_air_time) + self.last_contact_time = torch.zeros_like(self.last_air_time) + self.current_contact_time = torch.zeros_like(self.last_air_time) + + def reset(self, envs_idx: list[int] | None = None): + super().reset(envs_idx) + if envs_idx is None: + envs_idx = torch.arange(self.env.num_envs, device=gs.device) + + if not self.enabled: + return + + # reset the current air time + if self._track_air_time: + self.current_air_time[envs_idx] = 0.0 + self.current_contact_time[envs_idx] = 0.0 + self.last_air_time[envs_idx] = 0.0 + self.last_contact_time[envs_idx] = 0.0 + + def step(self): + super().step() + if not self.enabled: + return + self._calculate_contact_forces() + self._calculate_air_time() + + """ + Internal Implementation + """ + + def _get_links_idx( + self, entity_attr: str, names: list[str] = None + ) -> (torch.Tensor, torch.Tensor): + """ + Find the link indices for the given link names or regular expressions. + + Args: + entity: The entity to find the links in. + names: The names, or name regex patterns, of the links to find. + include_local_idx: Include a tensor of the local link indices, as well + + Returns: Tuple of global and local link index tensors. + """ + entity = self.env.__getattribute__(entity_attr) + + ids = [] + local_ids = [] + + if names is None: + # If link names are not defined, assume all links + for link in entity.links: + ids.append(link.idx) + local_ids.append(link.idx_local) + else: + for pattern in names: + found = False + for link in entity.links: + if pattern == link.name or re.match(f"^{pattern}$", link.name): + ids.append(link.idx) + local_ids.append(link.idx_local) + found = True + if not found: + names = [link.name for link in entity.links] + raise RuntimeError( + f"Link '{pattern}' not found in entity '{self._entity_attr}'.\nAvailable links: {names}" + ) + + return ( + torch.tensor(ids, device=gs.device), + torch.tensor(local_ids, device=gs.device), + ) + + def _calculate_contact_forces(self): + """ + Calculate contact forces using on the target links. + + Returns: + Tensor of shape (n_envs, n_target_links, 3) + """ + contacts = self.env.scene.rigid_solver.collider.get_contacts( + as_tensor=True, to_torch=True + ) + force = contacts["force"] + link_a = contacts["link_a"] + link_b = contacts["link_b"] + position = contacts["position"] + + # Validate physics engine outputs to prevent NaN/inf propagation + # Replace invalid values with zeros + if torch.isnan(force).any() or torch.isinf(force).any(): + force = torch.nan_to_num(force, nan=0.0, posinf=0.0, neginf=0.0) + print("Warning: Invalid contact forces detected (NaN/inf) and sanitized") + + # Get link quaternions used to transform the contact forces and positions into the local frame + links_quat = self.env.scene.rigid_solver.get_links_quat() + + # Clear output tensors + self.contacts.fill_(0.0) + self.contact_positions.fill_(0.0) + self._contact_position_counts.fill_(0.0) + + # Call unified kernel + kernel_get_contact_forces( + force.contiguous(), + position.contiguous(), + link_a.contiguous(), + link_b.contiguous(), + links_quat.contiguous(), + self._link_ids.contiguous(), + self._with_link_ids.contiguous(), + self.contacts.contiguous(), + self.contact_positions.contiguous(), + self._contact_position_counts.contiguous(), + 1 if self._has_with_filter else 0, + ) + + # Handle debug visualization + if self.debug_visualizer: + self._render_debug_visualizer( + self.contacts.clone().detach(), self.contact_positions.clone().detach() + ) + + def _calculate_air_time(self): + """ + Track air time values for the links + """ + if not self._track_air_time: + return + + dt = self.env.scene.dt + + # Check contact state of bodies + is_contact = ( + torch.norm(self.contacts[:, :, :], dim=-1) + > self._air_time_contact_threshold + ) + is_new_contact = (self.current_air_time > 0) * is_contact + is_new_detached = (self.current_contact_time > 0) * ~is_contact + + # Update the last contact time if body has just become in contact + self.last_air_time = torch.where( + is_new_contact, + self.current_air_time + dt, + self.last_air_time, + ) + + # Increment time for bodies that are not in contact + self.current_air_time = torch.where( + ~is_contact, + self.current_air_time + dt, + 0.0, + ) + + # Update the last contact time if body has just detached + self.last_contact_time = torch.where( + is_new_detached, + self.current_contact_time + dt, + self.last_contact_time, + ) + + # Increment time for bodies that are in contact + self.current_contact_time = torch.where( + is_contact, + self.current_contact_time + dt, + 0.0, + ) + + def _render_debug_visualizer( + self, contacts: torch.Tensor, contact_pos: torch.Tensor + ): + """ + Visualize contact points + + Args: + contacts: The contact forces experienced by the entity links. + contact_pos: The contact positions for each target link. + """ + # Clear existing debug objects + for node in self._debug_nodes: + self.env.scene.clear_debug_object(node) + self._debug_nodes = [] + + if not self.debug_visualizer: + return + + cfg = self.visualizer_cfg + + # Filter to only the environments we want to visualize + contacts = contacts[self.debug_envs_idx] + contact_pos = contact_pos[self.debug_envs_idx] + + # Filter out contacts below the force threshold + if "force_threshold" in cfg and cfg["force_threshold"] != 0.0: + force_mask = torch.norm(contacts, dim=-1) > cfg["force_threshold"] + contact_pos = contact_pos[force_mask] + + # Draw debug spheres + if contact_pos.numel() > 0: + node = self.env.scene.draw_debug_spheres( + poss=contact_pos, + radius=cfg["size"], + color=cfg["color"], + ) + if node is not None: + self._debug_nodes.append(node) + + def __repr__(self): + attrs = [f"link_names={self._link_names}"] + if self._entity_attr: + attrs.append(f"entity_attr={self._entity_attr}") + if self._with_entity_attr: + attrs.append(f"with_entity_attr={self._with_entity_attr}") + if self._with_links_names: + attrs.append(f"with_links_names={self._with_links_names}") + if self._track_air_time: + attrs.append(f"track_air_time={self._track_air_time}") + if self._air_time_contact_threshold: + attrs.append( + f"air_time_contact_threshold={self._air_time_contact_threshold}" + ) + attrs_str = ", ".join(attrs) + return f"{self.__class__.__name__}({attrs_str})" diff --git a/genesis_forge/managers/sensors/contact/kernel.py b/genesis_forge/managers/sensors/contact/kernel.py new file mode 100644 index 0000000..9424cbf --- /dev/null +++ b/genesis_forge/managers/sensors/contact/kernel.py @@ -0,0 +1,90 @@ +import gstaichi as ti +from genesis.utils.geom import ti_inv_transform_by_quat + + +@ti.kernel +def kernel_get_contact_forces( + contact_forces: ti.types.ndarray(), + contact_positions: ti.types.ndarray(), + link_a: ti.types.ndarray(), + link_b: ti.types.ndarray(), + links_quat: ti.types.ndarray(), + target_link_ids: ti.types.ndarray(), + with_link_ids: ti.types.ndarray(), + output_forces: ti.types.ndarray(), + output_positions: ti.types.ndarray(), + position_counts: ti.types.ndarray(), + has_with_filter: ti.i32, +): + """ + Accumulates contact forces and positions for target links, optionally filtering by with_link_ids. + + Args: + contact_forces: Contact force data (n_envs, n_contacts, 3) + contact_positions: Contact position data (n_envs, n_contacts, 3) + link_a: First link in each contact (n_envs, n_contacts) + link_b: Second link in each contact (n_envs, n_contacts) + links_quat: Link quaternions (n_envs, n_links, 4) + target_link_ids: Target link IDs to track (n_target_links) + with_link_ids: Filter links (n_with_links) - only used if has_with_filter=True + output_forces: Output force tensor (n_envs, n_target_links, 3) + output_positions: Output position tensor (n_envs, n_target_links, 3) + position_counts: Position count tensor (n_envs, n_target_links) - internal use only + has_with_filter: Whether to apply with_link filter (0 or 1) + """ + for i_b, i_c, i_t in ti.ndrange( + output_forces.shape[0], link_a.shape[-1], target_link_ids.shape[-1] + ): + contact_link_a = link_a[i_b, i_c] + contact_link_b = link_b[i_b, i_c] + target_link = target_link_ids[i_t] + + # Check if this contact involves our target link + is_target_a = contact_link_a == target_link + is_target_b = contact_link_b == target_link + + if is_target_a or is_target_b: + # Apply with_link filter if specified + should_include = True + if has_with_filter: + should_include = False + for i_w in range(with_link_ids.shape[-1]): + with_link = with_link_ids[i_w] + if (is_target_a and contact_link_b == with_link) or ( + is_target_b and contact_link_a == with_link + ): + should_include = True + break + + if should_include: + # Get contact force and position + force_vec = ti.Vector.zero(ti.f32, 3) + for j in ti.static(range(3)): + force_vec[j] = contact_forces[i_b, i_c, j] + output_positions[i_b, i_t, j] += contact_positions[i_b, i_c, j] + position_counts[i_b, i_t] += 1 + + # Get quaternions for both links + quat_a = ti.Vector.zero(ti.f32, 4) + quat_b = ti.Vector.zero(ti.f32, 4) + for j in ti.static(range(4)): + quat_a[j] = links_quat[i_b, contact_link_a, j] + quat_b[j] = links_quat[i_b, contact_link_b, j] + + # Transform force to local frame of target link + if is_target_b: + force_vec = ti_inv_transform_by_quat(force_vec, quat_b) + else: + force_vec = ti_inv_transform_by_quat(-force_vec, quat_a) + + # Accumulate force and position + for j in ti.static(range(3)): + output_forces[i_b, i_t, j] += force_vec[j] + + # Final pass: compute average positions for all links + for i_b, i_t in ti.ndrange(output_forces.shape[0], output_forces.shape[1]): + if position_counts[i_b, i_t] > 0: + for j in ti.static(range(3)): + output_positions[i_b, i_t, j] = ( + output_positions[i_b, i_t, j] / position_counts[i_b, i_t] + ) diff --git a/genesis_forge/managers/sensors/depth_camera_manager.py b/genesis_forge/managers/sensors/depth_camera_manager.py new file mode 100644 index 0000000..3f10263 --- /dev/null +++ b/genesis_forge/managers/sensors/depth_camera_manager.py @@ -0,0 +1,222 @@ +from __future__ import annotations + +import torch +import genesis as gs + +from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.managers.sensors.base_sensor_manager import BaseSensorManager +from typing import TYPE_CHECKING,Sequence + +if TYPE_CHECKING: + from genesis.engine.entities import RigidEntity + +class DepthCameraManager(BaseSensorManager): + """ + A Spherical Lidar attached to an entity's link in the environment. + + Args: + sensor_name: The name of the sensor + env: The environment to sense. + link_name: The name of the link the sensor is attached to + entity_attr: The environment attribute which contains the entity with the links we're tracking. Defaults to `robot`. + pos_offset: The xyz position offset between the link origin and the sensor origin + euler_offset: The xyz orintation offset between the link origin and the sensor origin + delay: The daly induced between sensor measurements + draw_debug: Boolean value indicating whether the sensors data should be visualised in the scene + read_frequency: the frequncy at which to read the input from the sensor in hz + res: tuple[int, int] + The resolution of the camera, specified as a tuple (width, height). + fx : float | None + Focal length in x direction in pixels. Computed from fov_horizontal if None. + fy : float | None + Focal length in y direction in pixels. Computed from fov_vertical if None. + cx : float | None + Principal point x coordinate in pixels. Defaults to image center if None. + cy : float | None + Principal point y coordinate in pixels. Defaults to image center if None. + fov_horizontal : float + Horizontal field of view in degrees. Used to compute fx if fx is None. + fov_vertical : float | None + Vertical field of view in degrees. Used to compute fy if fy is None. + angles: tuple[Sequence[float], Sequence[float]], optional + Array of horizontal/vertical angles. Overrides the other options if provided. + min_range : float, optional + The minimum sensing range in meters. Defaults to 0.0. + max_range : float, optional + The maximum sensing range in meters. Defaults to 20.0. + no_hit_value : float, optional + The value to return for no hit. Defaults to max_range if not specified. + return_world_frame : bool, optional + Whether to return points in the world frame. Defaults to False (local frame). + debug_sphere_radius: float, optional + The radius of each debug sphere drawn in the scene. Defaults to 0.02. + debug_ray_start_color: float, optional + The color of each debug ray start sphere drawn in the scene. Defaults to (0.5, 0.5, 1.0, 1.0). + debug_ray_hit_color: float, optional + The color of each debug ray hit point sphere drawn in the scene. Defaults to (1.0, 0.5, 0.5, 1.0). + """ + #TODO: write proper example + + def __init__( + self, + env: GenesisEnv, + sensor_name: str, + link_name: str, + entity_attr: RigidEntity = "robot", + pos_offset: list[float] = [0,0,0], + euler_offset: list[float] = [0,0,0], + delay: float = 0.0, + draw_debug: bool = False, + read_frequency: float= 25, + res: tuple[int, int]=None, + fx : float =None, + fy : float =None, + cx : float =None, + cy : float =None, + fov_horizontal : float=None, + fov_vertical : float = None, + min_range: float = 0.0, + max_range: float = 20.0, + no_hit_value: float = None, + return_world_frame: bool = False, + debug_sphere_radius: float = 0.02, + debug_ray_start_color: tuple[float, float, float, float] = (0.5, 0.5, 1.0, 1.0), + debug_ray_hit_color: tuple[float, float, float, float] = (1.0, 0.5, 0.5, 1.0), + read_rays=False, + read_images=False, + ): + super().__init__( + env=env, + sensor_name=sensor_name, + link_name=link_name, + entity_attr=entity_attr, + pos_offset=pos_offset, + euler_offset=euler_offset, + delay=delay, + draw_debug=draw_debug, + read_freuency=read_frequency) + + # scene=env.scene + depth_camera_pattern=gs.sensors.DepthCameraPattern( + res=res, + fx=fx,fy=fy, + cx=cx,cy=cy, + fov_horizontal=fov_horizontal, + fov_vertical=fov_vertical + ) + self._spherical_raycaster_sensor = self.env.scene.add_sensor( + gs.sensors.Lidar( + pattern=depth_camera_pattern, + entity_idx=self._entity.idx, + link_idx_local=self._link.idx_local, + pos_offset=self._pos_offset, + euler_offset=self._euler_offset, + delay=self._delay, + draw_debug=self._draw_debug, + min_range=min_range, + max_range=max_range, + no_hit_value=no_hit_value, + return_world_frame=return_world_frame, + debug_sphere_radius=debug_sphere_radius, + debug_ray_start_color=debug_ray_start_color, + debug_ray_hit_color=debug_ray_hit_color + ) + ) + self._read_rays=read_rays + self._read_images=read_images + + """ + Properties + """ + + @property + def link_name(self) -> torch.Tensor: + """name of the imu link.""" + return self._link_name + + @property + def link_idx(self) -> torch.Tensor: + """The link index for the imu link.""" + return self._link.idx + + @property + def local_link_idx(self) -> torch.Tensor: + """The local link index for the imu link.""" + return self._link.idx_local + + """ + Helper Methods + """ + def get_points(self) -> torch.Tensor: + """ + Get the point cloud measured by the sensor + + Returns: + The point cloud shape is (n_envs,n_points, 3) + """ + if self._read_rays: + return self._sensor_cloud[0] + gs.logger.error("the sensor was not initialised with the 'read_rays' arg, unable to get the points or distances") + return None + + def get_diatnaces(self) -> torch.Tensor: + """ + Get the distances measured by the sensor + + Returns: + The distances shape is (n_envs,n_points) + """ + if self._read_rays: + return self._sensor_cloud[1] + gs.logger.error("the sensor was not initialised with the 'read_rays' option, unable to get the points or distances") + return None + + def get_image(self)-> torch.Tensor: + """ + Get the depth image from the sensor + + Returns: + The distances shape is (n_envs,n_points) + """ + if self._read_images: + return self._sensor_image + gs.logger.error("the sensor was not initialised with the 'read_images' option, unable to get the depth image") + return None + + + """ + Lifecycle Operations + """ + + def build(self): + """Initialize link indices and buffers.""" + super().build() + self._num_points=self.sensor.read().shape[1] + if self._read_rays: + self._sensor_cloud=torch.zeros(self.env.num_envs,self._resolution[0],self._resolution[1],3) + if self._read_images: + self._sensor_image=torch.zeros(self.env.num_envs,self._resolution[0],self._resolution[1],3) + + def reset(self, envs_idx: list[int] | None = None): + super().reset(envs_idx) + if not self.enabled: + return + if self._read_rays: + self._sensor_cloud[self.envs_idx]=0.0 + if self._read_images: + self._sensor_image[self.envs_idx]=0.0 + + def step(self): + super().step() + if not self.enabled: + return + + if self._last_reading_timestamp is None: + self._last_reading_timestamp=self.env.scene.cur_t + elif self.env.scene.cur_t-self._last_reading_timestamp>self._sensor_read_interval: + if self._read_rays: + self._sensor_cloud[:]=self._spherical_raycaster_sensor.read() + if self._read_images: + self._sensor_cloud[:]=self._spherical_raycaster_sensor.read_image() + + self._last_reading_timestamp=self.env.scene.cur_t diff --git a/genesis_forge/managers/sensors/grid_raycaster_manager.py b/genesis_forge/managers/sensors/grid_raycaster_manager.py new file mode 100644 index 0000000..d4db8a8 --- /dev/null +++ b/genesis_forge/managers/sensors/grid_raycaster_manager.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +import torch +import genesis as gs + +from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.managers.sensors.base_sensor_manager import BaseSensorManager +from typing import TYPE_CHECKING,Sequence + +if TYPE_CHECKING: + from genesis.engine.entities import RigidEntity + +class GridRaycasterManager(BaseSensorManager): + """ + A grid Lidar attached to an entity's link in the environment. + + Args: + sensor_name: The name of the sensor + env: The environment to sense. + link_name: The name of the link the sensor is attached to + entity_attr: The environment attribute which contains the entity with the links we're tracking. Defaults to `robot`. + pos_offset: The xyz position offset between the link origin and the sensor origin + euler_offset: The xyz orintation offset between the link origin and the sensor origin + delay: The daly induced between sensor measurements + draw_debug: Boolean value indicating whether the sensors data should be visualised in the scene + read_frequency: the frequncy at which to read the input from the sensor in hz + resolution : float + Grid spacing in meters. + size : tuple[float, float] + Grid dimensions (length, width) in meters. + direction : tuple[float, float, float] + Ray direction vector. + min_range : float, optional + The minimum sensing range in meters. Defaults to 0.0. + max_range : float, optional + The maximum sensing range in meters. Defaults to 20.0. + no_hit_value : float, optional + The value to return for no hit. Defaults to max_range if not specified. + return_world_frame : bool, optional + Whether to return points in the world frame. Defaults to False (local frame). + debug_sphere_radius: float, optional + The radius of each debug sphere drawn in the scene. Defaults to 0.02. + debug_ray_start_color: float, optional + The color of each debug ray start sphere drawn in the scene. Defaults to (0.5, 0.5, 1.0, 1.0). + debug_ray_hit_color: float, optional + The color of each debug ray hit point sphere drawn in the scene. Defaults to (1.0, 0.5, 0.5, 1.0). + """ + #TODO: write proper example + + def __init__( + self, + env: GenesisEnv, + sensor_name: str, + link_name: str, + entity_attr: RigidEntity = "robot", + pos_offset: list[float] = [0,0,0], + euler_offset: list[float] = [0,0,0], + delay: float = 0.0, + draw_debug: bool = False, + read_frequency: float= 25, + resolution : float=0.1, + size : tuple[float, float]=(1,1), + direction : tuple[float, float, float]=(0,0,1), + min_range: float = 0.0, + max_range: float = 20.0, + no_hit_value: float = None, + return_world_frame: bool = False, + debug_sphere_radius: float = 0.02, + debug_ray_start_color: tuple[float, float, float, float] = (0.5, 0.5, 1.0, 1.0), + debug_ray_hit_color: tuple[float, float, float, float] = (1.0, 0.5, 0.5, 1.0), + ): + super().__init__( + env=env, + sensor_name=sensor_name, + link_name=link_name, + entity_attr=entity_attr, + pos_offset=pos_offset, + euler_offset=euler_offset, + delay=delay, + draw_debug=draw_debug, + read_freuency=read_frequency) + + # scene=env.scene + grid_pattern=gs.sensors.GridPattern( + resolution=resolution, + size=size, + direction=direction, + ) + self._spherical_raycaster_sensor = self.env.scene.add_sensor( + gs.sensors.Lidar( + pattern=grid_pattern, + entity_idx=self._entity.idx, + link_idx_local=self._link.idx_local, + pos_offset=self._pos_offset, + euler_offset=self._euler_offset, + delay=self._delay, + draw_debug=self._draw_debug, + min_range=min_range, + max_range=max_range, + no_hit_value=no_hit_value, + return_world_frame=return_world_frame, + debug_sphere_radius=debug_sphere_radius, + debug_ray_start_color=debug_ray_start_color, + debug_ray_hit_color=debug_ray_hit_color + ) + ) + + """ + Properties + """ + + @property + def link_name(self) -> torch.Tensor: + """name of the imu link.""" + return self._link_name + + @property + def link_idx(self) -> torch.Tensor: + """The link index for the imu link.""" + return self._link.idx + + @property + def local_link_idx(self) -> torch.Tensor: + """The local link index for the imu link.""" + return self._link.idx_local + + """ + Helper Methods + """ + def get_points(self) -> torch.Tensor: + """ + Get the point cloud measured by the sensor + + Returns: + The point cloud shape is (n_envs,n_points, 3) + """ + return self._sensor_reading[0] + + def get_diatnaces(self) -> torch.Tensor: + """ + Get the distances measured by the sensor + + Returns: + The distances shape is (n_envs,n_points) + """ + return self._sensor_reading[0] + + + """ + Lifecycle Operations + """ + + def build(self): + """Initialize link indices and buffers.""" + super().build() + self._num_points=self.sensor.read().shape[1] + self._sensor_reading=torch.zeros(self.env.num_envs,self._num_points,3) + + def reset(self, envs_idx: list[int] | None = None): + super().reset(envs_idx) + if not self.enabled: + return + + self._sensor_reading[self.envs_idx,:,:]=0.0 + + def step(self): + super().step() + if not self.enabled: + return + + if self._last_reading_timestamp is None: + self._last_reading_timestamp=self.env.scene.cur_t + elif self.env.scene.cur_t-self._last_reading_timestamp>self._sensor_read_interval: + self._sensor_reading[:]=self._spherical_raycaster_sensor.read() + self._last_reading_timestamp=self.env.scene.cur_t diff --git a/genesis_forge/managers/sensors/imu_manager.py b/genesis_forge/managers/sensors/imu_manager.py new file mode 100644 index 0000000..9b7430c --- /dev/null +++ b/genesis_forge/managers/sensors/imu_manager.py @@ -0,0 +1,230 @@ +from __future__ import annotations + +import torch +import genesis as gs + +from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.managers.sensors.base_sensor_manager import BaseSensorManager +from typing import TYPE_CHECKING,Sequence + +if TYPE_CHECKING: + from genesis.engine.entities import RigidEntity + +class ImuManager(BaseSensorManager): + """ + Inertal Measurenment Unit(IMU) sensor used to measure the lin_acc and ang_vel, it is attached to an entity's link in the environment. + + Args: + sensor_name: The name of the sensor + env: The environment to sense. + link_name: The name of the link the sensor is attached to + entity_attr: The environment attribute which contains the entity with the links we're tracking. Defaults to `robot`. + pos_offset: The xyz position offset between the link origin and the sensor origin + euler_offset: The xyz orintation offset between the link origin and the sensor origin + delay: The daly induced between sensor measurements + draw_debug: Boolean value indicating whether the sensors data should be visualised in the scene + read_frequency: the frequncy at which to read the input from the sensor in hz + queue_length: The length of the IMU reading queue buffer + acc_resolution : float, optional + The measurement resolution of the accelerometer (smallest increment of change in the sensor reading). + Default is 0.0, which means no quantization is applied. + acc_axes_skew : float | tuple[float, float, float] | Sequence[float] + Accelerometer axes alignment as a 3x3 rotation matrix, where diagonal elements represent alignment (0.0 to 1.0) + for each axis, and off-diagonal elements account for cross-axis misalignment effects. + - If a scalar is provided (float), all off-diagonal elements are set to the scalar value. + - If a 3-element vector is provided (tuple[float, float, float]), off-diagonal elements are set. + - If a full 3x3 matrix is provided, it is used directly. + acc_bias : tuple[float, float, float] + The constant additive bias for each axis of the accelerometer. + acc_noise : tuple[float, float, float] + The standard deviation of the white noise for each axis of the accelerometer. + acc_random_walk : tuple[float, float, float] + The standard deviation of the random walk, which acts as accumulated bias drift. + gyro_resolution : float, optional + The measurement resolution of the gyroscope (smallest increment of change in the sensor reading). + Default is 0.0, which means no quantization is applied. + gyro_axes_skew : float | tuple[float, float, float] | Sequence[float] + Gyroscope axes alignment as a 3x3 rotation matrix, similar to `acc_axes_skew`. + gyro_bias : tuple[float, float, float] + The constant additive bias for each axis of the gyroscope. + gyro_noise : tuple[float, float, float] + The standard deviation of the white noise for each axis of the gyroscope. + gyro_random_walk : tuple[float, float, float] + The standard deviation of the bias drift for each axis of the gyroscope. + debug_acc_color : float, optional + The rgba color of the debug acceleration arrow. Defaults to (0.0, 1.0, 1.0, 0.5). + debug_acc_scale: float, optional + The scale factor for the debug acceleration arrow. Defaults to 0.01. + debug_gyro_color : float, optional + The rgba color of the debug gyroscope arrow. Defaults to (1.0, 1.0, 0.0, 0.5). + debug_gyro_scale: float, optional + The scale factor for the debug gyroscope arrow. Defaults to 0.01. + + """ + #TODO: write proper example + + def __init__( + self, + env: GenesisEnv, + sensor_name: str, + link_name: str, + entity_attr: RigidEntity = "robot", + pos_offset: list[float] = [0,0,0], + euler_offset: list[float] = [0,0,0], + delay: float=0.0, + draw_debug: bool = False, + read_frequency: float=100, + queue_length: int = 25, + acc_resolution: float = 0.0, + acc_axes_skew: float | tuple[float, float, float] | Sequence[float] = 0.0, + acc_bias: tuple[float, float, float] = (0.0, 0.0, 0.0), + acc_noise: tuple[float, float, float] = (0.0, 0.0, 0.0), + acc_random_walk: tuple[float, float, float] = (0.0, 0.0, 0.0), + gyro_resolution: float = 0.0, + gyro_axes_skew: float | tuple[float, float, float] | Sequence[float] = 0.0, + gyro_bias: tuple[float, float, float] = (0.0, 0.0, 0.0), + gyro_noise: tuple[float, float, float] = (0.0, 0.0, 0.0), + gyro_random_walk: tuple[float, float, float] = (0.0, 0.0, 0.0), + debug_acc_color: tuple[float, float, float, float] = (0.0, 1.0, 1.0, 0.5), + debug_acc_scale: float = 0.01, + debug_gyro_color: tuple[float, float, float, float] = (1.0, 1.0, 0.0, 0.5), + debug_gyro_scale: float = 0.01, + ): + super().__init__( + env=env, + sensor_name=sensor_name, + link_name=link_name, + entity_attr=entity_attr, + pos_offset=pos_offset, + euler_offset=euler_offset, + delay=delay, + draw_debug=draw_debug, + read_freuency=read_frequency) + + self._queue_length = queue_length + + scene=env.scene + self._imu_sensor = scene.add_sensor( + gs.sensors.IMU( + **self.base_sensor_args, + acc_resolution = acc_resolution, + acc_axes_skew = acc_axes_skew, + acc_bias = acc_bias, + acc_noise = acc_noise, + acc_random_walk = acc_random_walk, + gyro_resolution = gyro_resolution, + gyro_axes_skew = gyro_axes_skew, + gyro_bias = gyro_bias, + gyro_noise = gyro_noise, + gyro_random_walk = gyro_random_walk, + debug_acc_color = debug_acc_color, + debug_acc_scale = debug_acc_scale, + debug_gyro_color = debug_gyro_color, + debug_gyro_scale = debug_gyro_scale, + ) + ) + + """ + Properties + """ + + @property + def link_name(self) -> torch.Tensor: + """name of the imu link.""" + return self._link_name + + @property + def link_idx(self) -> torch.Tensor: + """The link index for the imu link.""" + return self._link.idx + + @property + def local_link_idx(self) -> torch.Tensor: + """The local link index for the imu link.""" + return self._link.idx_local + + """ + Helper Methods + """ + def get_lin_acceleration(self) -> torch.Tensor: + """ + Get the linear acceleration of the imu + + Returns: + The linear acceleration of the imu. Shape is (n_envs, 3) + """ + return self._imu_reading[:,0,:] + + def get_ang_velocity(self) -> torch.Tensor: + """ + Get the angular velocity of the imu + + Returns: + The angular velocity of the imu. Shape is (n_envs, 3) + """ + return self._imu_reading[:,1,:] + + def get_lin_acceleration_queue(self) -> torch.Tensor: + """ + Get the linear acceleration buffer of the imu + + Returns: + The linear acceleration buffer of the imu. Shape is (n_envs, 3) + """ + return self._imu_queue[:,:,0,:] + + def get_ang_velocity_queue(self) -> torch.Tensor: + """ + Get the angular velocity buffer of the imu + + Returns: + The angular velocity buffer of the imu. Shape is (n_envs, 3) + """ + return self._imu_queue[:,:,1,:] + + """ + Lifecycle Operations + """ + + def build(self): + """Initialize buffers.""" + super().build() + + self._imu_reading=torch.zeros(self.num_envs,2,3) + self._imu_queue=torch.zeros(self.queue_length,self.num_envs,2,3) + + def reset(self, envs_idx: list[int] | None = None): + super().reset(envs_idx) + self._imu_reading[self.envs_idx,:,:]=0.0 + self._imu_queue[:,self.envs_idx,:,:]=0.0 + + + def step(self): + super().step() + imu_data=self._imu_sensor.read() + self._imu_reading[:,0,:]=imu_data[0] + self._imu_reading[:,1,:]=imu_data[1] + self._enqueue(self._imu_reading) + + """ + Internal Implementation + """ + def _enqueue(self,imu_reading: torch.Tensor): + if self._imu_queue.size(0) == self._queue_length: + self._imu_queue = self._imu_queue[1:] + self._imu_queue = torch.cat((self._imu_queue, imu_reading.unsqueeze(0))) + else: + self._imu_queue = torch.cat((self._imu_queue, imu_reading.unsqueeze(0))) + + def __repr__(self): + attrs = [f"link_name={self._link_name}"] + if self._entity_attr: + attrs.append(f"entity_attr={self._entity_attr}") + if self._pos_offset: + attrs.append(f"pos_offset={self._pos_offset}") + if self._euler_offset: + attrs.append(f"pos_offset={self._euler_offset}") + if self._queue_length: + attrs.append(f"queue_length={self._queue_length}") + attrs_str = ", ".join(attrs) + return f"{self.__class__.__name__}({attrs_str})" diff --git a/genesis_forge/managers/sensors/spherical_raycaster_manager.py b/genesis_forge/managers/sensors/spherical_raycaster_manager.py new file mode 100644 index 0000000..6287a8a --- /dev/null +++ b/genesis_forge/managers/sensors/spherical_raycaster_manager.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import torch +import genesis as gs + +from genesis_forge.genesis_env import GenesisEnv +from genesis_forge.managers.sensors.base_sensor_manager import BaseSensorManager +from typing import TYPE_CHECKING,Sequence + +if TYPE_CHECKING: + from genesis.engine.entities import RigidEntity + +class SphericalRaycasterManager(BaseSensorManager): + """ + A Spherical Lidar attached to an entity's link in the environment. + + Args: + sensor_name: The name of the sensor + env: The environment to sense. + link_name: The name of the link the sensor is attached to + entity_attr: The environment attribute which contains the entity with the links we're tracking. Defaults to `robot`. + pos_offset: The xyz position offset between the link origin and the sensor origin + euler_offset: The xyz orintation offset between the link origin and the sensor origin + delay: The daly induced between sensor measurements + draw_debug: Boolean value indicating whether the sensors data should be visualised in the scene + read_frequency: the frequncy at which to read the input from the sensor in hz + fov: tuple[float | tuple[float, float], float | tuple[float, float]] + Field of view in degrees for horizontal and vertical directions. Defaults to (360.0, 30.0). + If a single float is provided, the FOV is centered around 0 degrees. + If a tuple is provided, it specifies the (min, max) angles. + n_points: tuple[int, int] + Number of horizontal/azimuth and vertical/elevation scan lines. Defaults to (64, 128). + angular_resolution: tuple[float, float], optional + Horizontal and vertical angular resolution in degrees. Overrides n_points if provided. + angles: tuple[Sequence[float], Sequence[float]], optional + Array of horizontal/vertical angles. Overrides the other options if provided. + min_range : float, optional + The minimum sensing range in meters. Defaults to 0.0. + max_range : float, optional + The maximum sensing range in meters. Defaults to 20.0. + no_hit_value : float, optional + The value to return for no hit. Defaults to max_range if not specified. + return_world_frame : bool, optional + Whether to return points in the world frame. Defaults to False (local frame). + debug_sphere_radius: float, optional + The radius of each debug sphere drawn in the scene. Defaults to 0.02. + debug_ray_start_color: float, optional + The color of each debug ray start sphere drawn in the scene. Defaults to (0.5, 0.5, 1.0, 1.0). + debug_ray_hit_color: float, optional + The color of each debug ray hit point sphere drawn in the scene. Defaults to (1.0, 0.5, 0.5, 1.0). + """ + #TODO: write proper example + + def __init__( + self, + env: GenesisEnv, + sensor_name: str, + link_name: str, + entity_attr: RigidEntity = "robot", + pos_offset: list[float] = [0,0,0], + euler_offset: list[float] = [0,0,0], + delay: float = 0.0, + draw_debug: bool = False, + read_frequency: float= 25, + fov: tuple[float | tuple[float, float], float | tuple[float, float]]=(360.0, 30.0), + n_points: tuple[int, int]=(64, 128), + angular_resolution: tuple[float, float]=None, + angles: tuple[Sequence[float], Sequence[float]]=None, + min_range: float = 0.0, + max_range: float = 20.0, + no_hit_value: float = None, + return_world_frame: bool = False, + debug_sphere_radius: float = 0.02, + debug_ray_start_color: tuple[float, float, float, float] = (0.5, 0.5, 1.0, 1.0), + debug_ray_hit_color: tuple[float, float, float, float] = (1.0, 0.5, 0.5, 1.0), + ): + super().__init__( + env=env, + sensor_name=sensor_name, + link_name=link_name, + entity_attr=entity_attr, + pos_offset=pos_offset, + euler_offset=euler_offset, + delay=delay, + draw_debug=draw_debug, + read_freuency=read_frequency) + + # scene=env.scene + spherical_pattern=gs.sensors.SphericalPattern( + fov=fov, + n_points=n_points, + angular_resolution=angular_resolution, + angles=angles + ) + self._spherical_raycaster_sensor = self.env.scene.add_sensor( + gs.sensors.Lidar( + pattern=spherical_pattern, + entity_idx=self._entity.idx, + link_idx_local=self._link.idx_local, + pos_offset=self._pos_offset, + euler_offset=self._euler_offset, + delay=self._delay, + draw_debug=self._draw_debug, + min_range=min_range, + max_range=max_range, + no_hit_value=no_hit_value, + return_world_frame=return_world_frame, + debug_sphere_radius=debug_sphere_radius, + debug_ray_start_color=debug_ray_start_color, + debug_ray_hit_color=debug_ray_hit_color + ) + ) + + """ + Properties + """ + + @property + def link_name(self) -> torch.Tensor: + """name of the imu link.""" + return self._link_name + + @property + def link_idx(self) -> torch.Tensor: + """The link index for the imu link.""" + return self._link.idx + + @property + def local_link_idx(self) -> torch.Tensor: + """The local link index for the imu link.""" + return self._link.idx_local + + """ + Helper Methods + """ + def get_points(self) -> torch.Tensor: + """ + Get the point cloud measured by the sensor + + Returns: + The point cloud shape is (n_envs,n_points, 3) + """ + return self._sensor_reading[0] + + def get_diatnaces(self) -> torch.Tensor: + """ + Get the distances measured by the sensor + + Returns: + The distances shape is (n_envs,n_points) + """ + return self._sensor_reading[0] + + + """ + Lifecycle Operations + """ + + def build(self): + """Initialize link indices and buffers.""" + super().build() + self._num_points=self.sensor.read().shape[1] + self._sensor_reading=torch.zeros(self.env.num_envs,self._num_points,3) + + def reset(self, envs_idx: list[int] | None = None): + super().reset(envs_idx) + if not self.enabled: + return + + self._sensor_reading[self.envs_idx,:,:]=0.0 + + def step(self): + super().step() + if not self.enabled: + return + + if self._last_reading_timestamp is None: + self._last_reading_timestamp=self.env.scene.cur_t + elif self.env.scene.cur_t-self._last_reading_timestamp>self._sensor_read_interval: + self._sensor_reading[:]=self._spherical_raycaster_sensor.read() + self._last_reading_timestamp=self.env.scene.cur_t diff --git a/genesis_forge/mdp/reset.py b/genesis_forge/mdp/reset.py index f3de2aa..1950720 100644 --- a/genesis_forge/mdp/reset.py +++ b/genesis_forge/mdp/reset.py @@ -290,3 +290,64 @@ def __call__( links_idx_local=self._links_idx_local, envs_idx=envs_idx, ) + +class randomize_link_com_shift(ResetMdpFnClass): + """ + Randomly modify the center of mass(COM) of one or more links of the entity. + This picks a random shift from `x_range','y_range','z_range' bounds and passes it to `set_com_shift` for each environment. + + See: https://genesis-world.readthedocs.io/en/latest/api_reference/entity/rigid_entity/rigid_entity.html#genesis.engine.entities.rigid_entity.rigid_entity.RigidEntity.set_mass_shift + + Args: + env: The environment + entity: The entity to set the rotation of. + link_name: The name, or regex pattern, of the link(s) to set the mass for. + com_shift_ranges: The range of distance to add to the center of mass(mass) + """ + + def __init__( + self, + _env: GenesisEnv, + entity: RigidEntity, + link_name: str, + com_shift_ranges: tuple[tuple[float, float],tuple[float, float],tuple[float, float]], + ): + self.env = _env + self._entity = entity + self._link_name = link_name + self._links_idx_local = [] + self._com_shift_buffer: torch.tensor | None = None + self.build() + + def build(self): + self._links_idx_local = [] + self._orig_mass = None + if self._link_name is not None: + links = links_by_name_pattern(self._entity, self._link_name) + if len(links) > 0: + self._links_idx_local = [link.idx_local for link in links] + self._mass_shift_buffer = torch.zeros( + (self.env.num_envs, len(self._links_idx_local),3), device=gs.device + ) + else: + raise ValueError( + f"No links found with name/pattern '{self._link_name}'" + ) + + def __call__( + self, + env: GenesisEnv, + entity: RigidEntity, + envs_idx: list[int], + link_name: str, + com_shift_ranges: tuple[tuple[float, float],tuple[float, float],tuple[float, float]], + ): + # Randomize mass + self._mass_shift_buffer[envs_idx, :].uniform_(*com_shift_ranges) + + # Set mass on entity + self._entity.set_mass_shift( + self._mass_shift_buffer[envs_idx], + links_idx_local=self._links_idx_local, + envs_idx=envs_idx, + ) diff --git a/genesis_forge/mdp/rewards.py b/genesis_forge/mdp/rewards.py index eae526b..d471044 100644 --- a/genesis_forge/mdp/rewards.py +++ b/genesis_forge/mdp/rewards.py @@ -7,13 +7,17 @@ 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, + ImuManager, TerrainManager, EntityManager, ) @@ -52,7 +56,6 @@ def terminated(env: GenesisEnv) -> torch.Tensor: Robot base position/state """ - def base_height( env: GenesisEnv, target_height: Union[float, torch.Tensor] = None, @@ -283,9 +286,213 @@ def action_rate_l2(env: GenesisEnv) -> torch.Tensor: """ -Velocity Command Rewards +Position Command Rewards """ +def commonad_tracking_base_position( + env: GenesisEnv, + command: torch.Tensor = None, + position_cmd_manager: PositionCommandManager = None, + sensitivity: float = 0.25, + entity_attr: str = "robot", + entity_manager: EntityManager = None, +) -> torch.Tensor: + """ + Penalize base pose 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 velocity 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: + base_pos = entity_manager.entity.get_pos() + else: + robot = getattr(env, entity_attr) + base_pos = robot.get_pos() + + if position_cmd_manager is not None: + command = position_cmd_manager.command[:, :2] + + base_pos_error= torch.sum(torch.square(command - base_pos), dim=1) + return torch.square(-base_pos_error/sensitivity) + +def commonad_tracking_link_position( + env: GenesisEnv, + link_name: str=None, + command: torch.Tensor = None, + position_cmd_manager: PositionCommandManager = None, + sensitivity: float = 0.25, + entity_attr: str = "robot", + entity_manager: EntityManager = None, +) -> torch.Tensor: + """ + Penalize link pose away from target. + + Args: + env: The Genesis environment containing the robot + link_name: the name of the position to track + command: The commanded XYZ position the the world frame, its shape is(num_envs, 3) + position_cmd_manager: The velocity 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: + link_pos = entity_manager.entity.get_link(link_name).get_pos() + else: + robot = getattr(env, entity_attr) + link_pos = robot.get_link(link_name).get_pos() + + if position_cmd_manager is not None: + command = position_cmd_manager.command[:, :2] + + link_pos_error= torch.sum(torch.square(command - link_pos), dim=1) + return torch.square(-link_pos_error/sensitivity) + +""" +Pose Command Rewards +""" + +def command_tracking_base_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", + entity_manager: 'EntityManager' = None, +) -> torch.Tensor: + """ + Penalize base pose away from target (both position and Euler angles) 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. + position_cmd_manager: The velocity 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 position_cmd_manager must be provided to command_tracking_base_pose" + + if entity_manager is not None: + base_pos = entity_manager.entity.get_pos() + base_euler = quat_to_xyz(entity_manager.entity.get_euler()) + else: + robot = getattr(env, entity_attr) + base_pos = robot.get_pos() + base_euler = quat_to_xyz(robot.get_euler()) + + if pose_cmd_manager is not None: + command_pos = pose_cmd_manager.command[:, :3] + command_euler = pose_cmd_manager.command[:, 3:6] + else: + command_pos = command[:, :3] + command_euler = command[:, 3:6] + + pos_error = torch.sum(torch.square(command_pos - base_pos), dim=1) + + euler_error = torch.sum(torch.square( + torch.min( + torch.abs(command_euler - base_euler), + 2 * torch.pi - torch.abs(command_euler - base_euler) + ) + ), dim=1) + + pos_penalty = torch.square(-pos_error / pos_sensitivity) + euler_penalty = torch.square(-euler_error / euler_sensitivity) + + total_penalty = pos_penalty + euler_penalty + + return total_penalty + +def command_tracking_link_pose( + env: GenesisEnv, + link_name: str, + command: torch.Tensor = None, + pose_cmd_manager: PoseCommandManager = None, + pos_sensitivity: float = 0.25, + euler_sensitivity: float = 0.25, + entity_attr: str = "robot", + entity_manager: 'EntityManager' = None, +) -> torch.Tensor: + """ + Penalize base pose away from target (both position and Euler angles) 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. + position_cmd_manager: The velocity 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 position_cmd_manager must be provided to command_tracking_base_pose" + + if entity_manager is not None: + base_pos = entity_manager.entity.get_link(link_name).get_pos() + base_euler = quat_to_xyz(entity_manager.entity.get_link(link_name).get_euler()) + else: + robot = getattr(env, entity_attr) + base_pos = robot.get_link(link_name).get_pos() + base_euler = quat_to_xyz(robot.get_link(link_name).get_euler()) + + if pose_cmd_manager is not None: + command_pos = pose_cmd_manager.command[:, :3] + command_euler = pose_cmd_manager.command[:, 3:6] + else: + command_pos = command[:, :3] + command_euler = command[:, 3:6] + + pos_error = torch.sum(torch.square(command_pos - base_pos), dim=1) + + euler_error = torch.sum(torch.square( + torch.min( + torch.abs(command_euler - base_euler), + 2 * torch.pi - torch.abs(command_euler - base_euler) + ) + ), dim=1) + + pos_penalty = torch.square(-pos_error / pos_sensitivity) + euler_penalty = torch.square(-euler_error / euler_sensitivity) + + total_penalty = pos_penalty + euler_penalty + + return total_penalty + +""" +Velocity Command Rewards +""" def command_tracking_lin_vel( env: GenesisEnv, @@ -303,9 +510,9 @@ def command_tracking_lin_vel( command: The commanded XY linear velocity in the shape (num_envs, 2) vel_cmd_manager: The velocity command manager sensitivity: A lower value means the reward is more sensitive to the error - entity_manager: The entity manager for the robot/entity the reward is being computed for. - This is slightly more performant than using the `entity_attr` parameter. entity_attr: The attribute name of the entity in the environment. This isn't necessary if `entity_manager` is provided. + entity_manager: The entity manager for the robot/entity the reward is being computed for. + This is slightly more performant than using the `entity_attr` parameter. Returns: torch.Tensor: Reward for tracking of linear velocity commands (xy axes) @@ -406,6 +613,63 @@ def stand_still_joint_deviation_l1( return joint_deviation * (torch.norm(command[:, :2], dim=1) < command_threshold) +""" +Imu +""" + +def imu_lin_acc_jitter( + _env: GenesisEnv, + imu_manager: ImuManager, + queue_size: int = 10, + ignore_gravity: bool = False, +) -> torch.Tensor: + """ + Penalize changes in linear acceleration (jerk) using the IMU's internal queue. + + Returns: + torch.Tensor of shape (num_envs,) + """ + lin_acc_queue = imu_manager.get_lin_acceleration_queue() + T = lin_acc_queue.shape[0] + + if ignore_gravity: + lin_acc_queue[..., 2] += 9.81 + + if queue_size < T: + lin_acc_queue = lin_acc_queue[-queue_size:] + + if lin_acc_queue.shape[0] < 2: + return torch.zeros(_env.scene.num_envs) + + diffs = lin_acc_queue[1:] - lin_acc_queue[:-1] + mags = torch.norm(diffs, dim=-1) + return mags.sum(dim=0) + +def imu_ang_vel_jitter( + _env: GenesisEnv, + imu_manager: ImuManager, + queue_size: int = 10, +) -> torch.Tensor: + """ + Penalize changes in angular velocity (gyro jerk) using the IMU's internal queue. + + Returns: + torch.Tensor of shape (num_envs,) + """ + ang_vel_queue = imu_manager.get_ang_velocity_queue() + T = ang_vel_queue.shape[0] + + if queue_size < T: + ang_vel_queue = ang_vel_queue[-queue_size:] + + if ang_vel_queue.shape[0] < 2: + return torch.zeros(_env.scene.num_envs) + + diffs = ang_vel_queue[1:] - ang_vel_queue[:-1] + mags = torch.norm(diffs, dim=-1) + + return mags.sum(dim=0) + """ Contacts """ diff --git a/genesis_forge/ros2_env.py b/genesis_forge/ros2_env.py new file mode 100644 index 0000000..56e9057 --- /dev/null +++ b/genesis_forge/ros2_env.py @@ -0,0 +1,292 @@ +from __future__ import annotations +import math +import torch +from gymnasium import spaces +from typing import Any, Literal, TYPE_CHECKING + +if TYPE_CHECKING: + from genesis.engine.entities import RigidEntity + +EnvMode = Literal["train", "eval", "play"] + + +class Ros2Env: + """ + Base environment class for your simulated robot environment. + + Args: + num_envs: Number of parallel environments. + dt: Simulation time step. + max_episode_length_sec: Maximum episode length in seconds. + max_episode_random_scaling: Scale the maximum episode length by this amount (+/-) so that not all environments reset at the same time. + extras_logging_key: The key used, in info/extras dict, which is returned by step and reset functions, to send data to tensorboard by the RL agent. + + Example:: + + class MyEnv(GenesisEnv): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # ...Define scene here... + self.scene = gs.Scene() + self.terrain = self.scene.add_entity(gs.morphs.Plane()) + self.robot = self.scene.add_entity( ... ) + + def step(self, actions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: + # ...step logic here... + return obs, rewards, terminations, truncations, info + + def reset(self, envs_idx: list[int] = None) -> tuple[torch.Tensor, dict[str, Any]]: + # ...reset logic here... + return obs, info + + def get_observations(self) -> torch.Tensor: + # ...define current observations here... + return obs + + """ + + action_space: spaces.Space | None = None + observation_space: spaces.Space | None = None + can_be_wrapped: bool = True + + def __init__( + self, + num_envs: int = 1, + dt: float = 1 / 100, + max_episode_length_sec: int | None = 10, + max_episode_random_scaling: float = 0.0, + extras_logging_key: str = "episode", + ): + self.dt = dt + self.num_envs = num_envs + if torch.cuda.is_available(): + self.device = torch.device("cuda:0") + else: + self.device = torch.device("cpu") # Fall back to CPU if no GPU is available + + self.extras_logging_key = extras_logging_key + self._extras = {} + self._extras[extras_logging_key] = {} + + self._actions: torch.Tensor = None + self._last_actions: torch.Tensor = None + + self.step_count: int = 0 + self.episode_length = torch.zeros( + (self.num_envs,), device=self.device, dtype=torch.int32 + ) + self.max_episode_length: torch.Tensor = None + + self._max_episode_length_sec = 0.0 + self._base_max_episode_length = None + self._max_episode_random_scaling = max_episode_random_scaling + if max_episode_length_sec and max_episode_length_sec > 0: + self.max_episode_length = torch.zeros( + (self.num_envs,), device=self.device, dtype=torch.int + ) + self.max_episode_length[:] = self.set_max_episode_length( + max_episode_length_sec + ) + + """ + Properties + """ + + @property + def unwrapped(self): + """Returns this environment, not a wrapped version of it.""" + return self + + @property + def max_episode_length_sec(self) -> int | None: + """The max episode length, in seconds, for each environment.""" + return self._max_episode_length_sec + + @property + def extras(self) -> dict: + """ + The extras/infos dictionary reset at the start of every step, and contains additional data about the environment during that step. + """ + return self._extras + + @property + def actions(self) -> torch.Tensor: + """ + The actions for each environment for this step. + If you're using an action manager, these are the actions prior to being handled by the action manager. + """ + return self._actions + + @property + def last_actions(self) -> torch.Tensor: + """ + The actions for for the previous step. + """ + return self._last_actions + + @property + def num_actions(self) -> int: + """The number of actions for each environment.""" + if self.action_space is not None: + return self.action_space.shape[0] + return 0 + + @property + def num_observations(self) -> int: + """The number of observations for each environment.""" + if self.observation_space is not None: + return self.observation_space.shape[0] + return 0 + + @property + def max_episode_length_steps(self) -> int | None: + """ + The max episode length, in steps, for each environment. + If episode randomization scaling is enabled, this will be the base max episode length before scaling. + """ + return self._base_max_episode_length + + """ + Utilities + """ + + def set_max_episode_length(self, max_episode_length_sec: int) -> int: + """ + Set or change the maximum episode length. + + Args: + max_episode_length_sec: The maximum episode length in seconds. + + Returns: + The maximum episode length in steps. + """ + self._max_episode_length_sec = max_episode_length_sec + self._base_max_episode_length = math.ceil(max_episode_length_sec / self.dt) + return self._base_max_episode_length + + """ + Operations + """ + + def build(self) -> None: + """ + Builds the environment before the first step. + The Genesis scene and all the scene entities must be added before calling this method. + """ + assert ( + self.scene is not None + ), "The scene must be constructed and assigned to the .scene attribute before building." + self.scene.build(n_envs=self.num_envs) + + def step( + self, actions: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: + """ + Performs a step in all environments with the given actions. + + Args: + actions: Batch of actions for each environment with the :attr:`action_space` shape. + + Returns: + Batch of (observations, rewards, terminations, truncations, info/extras) + """ + self._extras = {} + self._extras[self.extras_logging_key] = {} + self.step_count += 1 + self.episode_length += 1 + + if self._actions is None: + self._actions = actions.detach().clone() + self._last_actions = torch.zeros_like(actions, device=self.device) + else: + self._last_actions[:] = self._actions[:] + self._actions[:] = actions[:] + + return None, None, None, None, self._extras + + def reset( + self, + envs_idx: list[int] = None, + ) -> tuple[torch.Tensor, dict[str, Any]]: + """ + Reset one or more environments. + Each of the registered managers will also be reset for those environments. + + Args: + env_ids: The environment ids to reset. If None, all environments are reset. + + Returns: + A batch of observations and info from the vectorized environment. + """ + if envs_idx is None: + envs_idx = torch.arange(self.num_envs, device=self.device) + + # Initial reset, set buffers + if self.step_count == 0 and self.action_space is not None: + self._actions = torch.zeros( + (self.num_envs, self.action_space.shape[0]), + device=self.device, + dtype=torch.float, + ) + self._last_actions = torch.zeros_like(self.actions, device=self.device) + + # Actions + if envs_idx.numel() > 0: + if self.actions is not None: + self.actions[envs_idx] = 0.0 + self._last_actions[envs_idx] = 0.0 + + # Episode length + self.episode_length[envs_idx] = 0 + + # Randomize max episode length for env_ids + if ( + len(envs_idx) > 0 + and self._max_episode_random_scaling > 0.0 + and self._base_max_episode_length is not None + ): + max_random_scaling = ( + self._base_max_episode_length * self._max_episode_random_scaling + ) + randomization = ( + torch.empty((envs_idx.numel(),)).uniform_(-1.0, 1.0) + * max_random_scaling + ) + self.max_episode_length[envs_idx] = torch.round( + self._base_max_episode_length + randomization + ).to(torch.int) + + return None, self.extras + + def get_observations(self) -> torch.Tensor: + """ + Returns the current observations for each environment. + Override this method to return the observations for your environment. + + Example:: + + def get_observations(self) -> torch.Tensor: + return torch.cat( + [ + self.base_ang_vel * self.obs_scales["ang_vel"], # 3 + self.projected_gravity, # 3 + self.commands * self.commands_scale, # 3 + (self.dof_pos - self.default_dof_pos) * self.obs_scales["dof_pos"], # 12 + self.dof_vel * self.obs_scales["dof_vel"], # 12 + self.actions, # 12 + ], + axis=-1, + ) + """ + if self.observation_space is not None: + return torch.zeros( + (self.num_envs, self.observation_space.shape[0]), + device=self.device, + dtype=torch.float, + ) + return None + + def close(self): + """Close the environment.""" + pass diff --git a/genesis_forge/ros2_managed_env.py b/genesis_forge/ros2_managed_env.py new file mode 100644 index 0000000..9dd5592 --- /dev/null +++ b/genesis_forge/ros2_managed_env.py @@ -0,0 +1,401 @@ +import torch +from typing import Any, TypedDict +from gymnasium import spaces +from tensordict import TensorDict +from genesis_forge.ros2_env import Ros2Env +from genesis_forge.managers.base import BaseManager, ManagerType +from genesis_forge.managers import ( + ContactManager, + EntityManager, + CommandManager, + TerrainManager, + PositionActionManager, + ObservationManager, + RewardManager, + TerminationManager, + ActuatorManager, +) + + +class ManagersDict(TypedDict): + actuator: ActuatorManager | None + contact: list[ContactManager] + entity: list[EntityManager] + command: list[CommandManager] + terrain: list[TerrainManager] + action: PositionActionManager | None + observation: list[ObservationManager] + reward: RewardManager | None + termination: TerminationManager | None + + +class Ros2ManagedEnvironment(Ros2Env): + """ + An environment which moves a lot of the logic of the environment to manager classes. + This helps to keep the environment code clean and modular. + + Args: + num_envs: Number of parallel environments. + dt: Simulation time step. + max_episode_length_sec: Maximum episode length in seconds. + max_episode_random_scaling: Randomly scale the maximum episode length by this amount (+/-) so that not all environments reset at the same time. + extras_logging_key: The key used, in info/extras dict, which is returned by step and reset functions, to send data to tensorboard by the RL agent. + + Example:: + + class MyEnv(ManagedEnvironment): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # ...Define scene here... + + def config(self): + self.action_manager = PositionalActionManager( + self, + joint_names=".*", + pd_kp=50, + pd_kv=0.5, + max_force=8.0, + default_pos={ + # Hip joints + "Leg[1-2]_Hip": -1.0, + "Leg[3-4]_Hip": 1.0, + # Femur joints + "Leg[1-4]_Femur": 0.5, + # Tibia joints + "Leg[1-4]_Tibia": 0.6, + }, + ) + self.reward_manager = RewardManager( + self, + term_cfg={ + "Default pose": { + "weight": -1.0, + "fn": rewards.dof_similar_to_default, + "params": { + "dof_action_manager": self.action_manager, + }, + }, + "Base height": { + "fn": mdp.rewards.base_height, + "params": { "target_height": 0.135 }, + "weight": -100.0, + }, + }, + ) + ObservationManager( + self, + cfg={ + "velocity_cmd": {"fn": self.velocity_command.observation}, + "robot_ang_vel": { + "fn": utils.entity_ang_vel, + "params": {"entity": self.robot}, + "noise": 0.1, + }, + "robot_lin_vel": { + "fn": utils.entity_lin_vel, + "params": {"entity": self.robot}, + "noise": 0.1, + }, + "robot_projected_gravity": { + "fn": utils.entity_projected_gravity, + "params": {"entity": self.robot}, + "noise": 0.1, + }, + "robot_dofs_position": { + "fn": self.action_manager.get_dofs_position, + "noise": 0.01, + }, + "actions": {"fn": lambda: env.actions}, + }, + ) + """ + + def __init__( + self, + num_envs: int = 1, + dt: float = 1 / 100, + max_episode_length_sec: int | None = 10, + max_episode_random_scaling: float = 0.0, + extras_logging_key: str = "episode", + ): + super().__init__( + num_envs=num_envs, + dt=dt, + max_episode_length_sec=max_episode_length_sec, + max_episode_random_scaling=max_episode_random_scaling, + extras_logging_key=extras_logging_key, + ) + self.managers: ManagersDict = { + "contact": [], + "entity": [], + "command": [], + "terrain": [], + # there can only be one of each of these + "actuator": None, + "action": None, + "observation": [], + "reward": None, + "termination": None, + } + + self._action_space = None + self._observation_space = None + self._reward_buf = torch.zeros( + (self.num_envs,), device=self.device, dtype=torch.float + ) + self._terminated_buf = torch.zeros( + (self.num_envs,), device=self.device, dtype=torch.bool + ) + self._truncated_buf = torch.zeros( + (self.num_envs,), device=self.device, dtype=torch.bool + ) + + """ + Properties + """ + + @property + def action_space(self) -> torch.Tensor: + """ + The action space, provided by the action manager, if it exists. + """ + if self.managers["action"] is not None: + return self.managers["action"].action_space + if self._action_space is not None: + return self._action_space + return None + + @action_space.setter + def action_space(self, action_space: spaces.Space): + """ + Set the action space. + """ + self._action_space = action_space + + @property + def observation_space(self) -> spaces.Space: + """ + The observation space for the "policy" observation manager, if it exists. + """ + if len(self.managers["observation"]) > 0: + for obs in self.managers["observation"]: + if obs.name == "policy": + return obs.observation_space + return self.managers["observation"][0].observation_space + if self._observation_space is not None: + return self._observation_space + return None + + @observation_space.setter + def observation_space(self, observation_space: spaces.Space): + """ + Set the observation space. + """ + self._observation_space = observation_space + + """ + Managers + """ + + def add_manager(self, manager_type: ManagerType, manager: BaseManager): + """ + Adds a manager to the environment. + This will automatically be called by the manager class. + + Args: + manager_type: The type of manager to add. + manager: The manager to add. + """ + if manager_type not in self.managers: + raise ValueError(f"'{manager_type}' is not a valid manager type.") + + # Append manager if the dict item is a list + if isinstance(self.managers[manager_type], list): + self.managers[manager_type].append(manager) + elif self.managers[manager_type] is None: + self.managers[manager_type] = manager + else: + raise ValueError( + f"Manager type '{manager_type}' already has a manager, and an environment cannot have multiple {manager_type} managers." + ) + + """ + Operations + """ + + def config(self): + """ + Override this method and initialize all your managers here. + + Example:: + + def config(self): + EntityManager( + self, + entity_attr="robot", + on_reset={ + "position": { + "fn": reset.position, + "params": { + "position": INITIAL_BODY_POSITION, + "quat": INITIAL_QUAT, + }, + }, + }, + ) + """ + pass + + def build(self): + """ + Builds the environment before the first step. + The Genesis scene and all the scene entities must be added before calling this method. + """ + super().build() + self.config() + + for terrain_manager in self.managers["terrain"]: + terrain_manager.build() + if self.managers["actuator"] is not None: + self.managers["actuator"].build() + if self.managers["action"] is not None: + self.managers["action"].build() + for contact_manager in self.managers["contact"]: + contact_manager.build() + if self.managers["termination"] is not None: + self.managers["termination"].build() + if self.managers["reward"] is not None: + self.managers["reward"].build() + for command_manager in self.managers["command"]: + command_manager.build() + for entity_manager in self.managers["entity"]: + entity_manager.build() + for obs in self.managers["observation"]: + obs.build() + + def step( + self, actions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: + """ + Performs a step in all environments with the given actions. + + Args: + actions: Batch of actions for each environment with the :attr:`action_space` shape. + + Returns: + Batch of (observations, rewards, terminations, truncations, extras) + """ + super().step(actions) + self.extras["observations"] = TensorDict({}, device=self.device) + + # Execute the actions and a simulation step + if self.managers["action"] is not None: + self.managers["action"].step(actions) + + # Update entity managers + for entity_manager in self.managers["entity"]: + entity_manager.step() + + # Calculate contact forces + for contact_manager in self.managers["contact"]: + contact_manager.step() + + # Calculate termination and truncation + reset_env_idx = None + truncated = self._truncated_buf + terminated = self._terminated_buf + if self.managers["termination"] is not None: + terminated, truncated = self.managers["termination"].step() + reset_env_idx = ( + (terminated | truncated).nonzero(as_tuple=False).reshape((-1,)).detach() + ) + + # Calculate rewards + rewards = self._reward_buf + if self.managers["reward"] is not None: + rewards = self.managers["reward"].step() + + # Command managers + for command_manager in self.managers["command"]: + command_manager.step() + + # Reset environments + if reset_env_idx is not None and reset_env_idx.numel() > 0: + self.reset(reset_env_idx) + + # Get observations + obs = self.get_observations() + return ( + obs, + rewards, + terminated, + truncated, + self.extras, + ) + + def reset( + self, env_ids: list[int] | None = None + ) -> tuple[torch.Tensor, dict[str, Any]]: + """ + Reset one or more environments. + Each of the registered managers will also be reset for those environments. + + Args: + env_ids: The environment ids to reset. If None, all environments are reset. + + Returns: + A batch of observations (if env_ids is None) and an info dictionary from the vectorized environment. + """ + (obs, _) = super().reset(env_ids) + + if self.managers["actuator"] is not None: + self.managers["actuator"].reset(env_ids) + if self.managers["action"] is not None: + self.managers["action"].reset(env_ids) + for entity_manager in self.managers["entity"]: + entity_manager.reset(env_ids) + for contact_manager in self.managers["contact"]: + contact_manager.reset(env_ids) + if self.managers["termination"] is not None: + self.managers["termination"].reset(env_ids) + if self.managers["reward"] is not None: + self.managers["reward"].reset(env_ids) + for command_manager in self.managers["command"]: + command_manager.reset(env_ids) + for obs_manager in self.managers["observation"]: + obs_manager.reset(env_ids) + + # Only get observations when env_ids is None because this will be the initial reset called before the first step + # Otherwise, the observations are ignored + if env_ids is None: + obs = self.get_observations() + + return obs, self.extras + + def get_observations(self) -> torch.Tensor: + """ + Returns the current observations for this step. + If you use the ObservationManager, this will be handled automatically. + Otherwise, override this method to return the observations. + """ + if len(self.managers["observation"]) > 0: + # We already have observations for this step + if ( + "observations" in self.extras + and "policy" in self.extras["observations"] + ): + return self.extras["observations"]["policy"] + if "observations" not in self.extras: + self.extras["observations"] = TensorDict({}, device=gs.device) + + # Get observations + policy_obs = None + for obs_manager in self.managers["observation"]: + obs = obs_manager.get_observations() + self.extras["observations"][obs_manager.name] = obs + if obs_manager.name == "policy": + policy_obs = obs + return policy_obs + + return super().get_observations()