From c8a7e2727108a3a449d643056ae8574b935c31c7 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:02:37 +0530 Subject: [PATCH 01/29] add Actuator manager import --- examples/rough_terrain/environment.py | 1 + 1 file changed, 1 insertion(+) 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 From 54fab99f9368998be6a0520ac11785c2a98a1324 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:05:35 +0530 Subject: [PATCH 02/29] add skrl library examples --- examples/skrl/README.md | 40 ++++++ examples/skrl/environment.py | 272 +++++++++++++++++++++++++++++++++++ examples/skrl/eval.py | 87 +++++++++++ examples/skrl/train.py | 158 ++++++++++++++++++++ 4 files changed, 557 insertions(+) create mode 100644 examples/skrl/README.md create mode 100644 examples/skrl/environment.py create mode 100644 examples/skrl/eval.py create mode 100644 examples/skrl/train.py 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() From cdffe8af93e8712a158bcfb6e5967024b32b14a7 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:07:32 +0530 Subject: [PATCH 03/29] add imu sensor manager --- genesis_forge/managers/sensors/imu_manager.py | 230 ++++++++++++++++++ 1 file changed, 230 insertions(+) create mode 100644 genesis_forge/managers/sensors/imu_manager.py 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})" From c672fdf07c499738ff3cab62b3aa4ae9b7e113d0 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:07:50 +0530 Subject: [PATCH 04/29] add camera sensor manager --- .../managers/sensors/camera_manager.py | 317 ++++++++++++++++++ 1 file changed, 317 insertions(+) create mode 100644 genesis_forge/managers/sensors/camera_manager.py 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 + ) + + From 7c172dfcddde31738187fa9f76b917e4f81aa5d9 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:08:15 +0530 Subject: [PATCH 05/29] add depth camera sensor manager --- .../managers/sensors/depth_camera_manager.py | 222 ++++++++++++++++++ 1 file changed, 222 insertions(+) create mode 100644 genesis_forge/managers/sensors/depth_camera_manager.py 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 From e80bc92aa978a1e92d7f9377b1b0a4555593ca09 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:08:43 +0530 Subject: [PATCH 06/29] add base sensor manager --- .../managers/sensors/base_sensor_manager.py | 157 ++++++++++++++++++ 1 file changed, 157 insertions(+) create mode 100644 genesis_forge/managers/sensors/base_sensor_manager.py 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})" From 0535518a16e0cc5818c6f64cb619d3cd6398ae17 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:09:17 +0530 Subject: [PATCH 07/29] add grid raycaster sensor manager --- .../sensors/grid_raycaster_manager.py | 175 ++++++++++++++++++ 1 file changed, 175 insertions(+) create mode 100644 genesis_forge/managers/sensors/grid_raycaster_manager.py 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 From 5e954f31762e76e190b811276badaf0cd0c9026f Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:09:36 +0530 Subject: [PATCH 08/29] add spherical raycaster sensor manager --- .../sensors/spherical_raycaster_manager.py | 181 ++++++++++++++++++ 1 file changed, 181 insertions(+) create mode 100644 genesis_forge/managers/sensors/spherical_raycaster_manager.py 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 From 8263767f483b39242bb5e9d6cbae6ef60302af86 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:10:07 +0530 Subject: [PATCH 09/29] add sensor managers init --- genesis_forge/managers/sensors/__init__.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 genesis_forge/managers/sensors/__init__.py 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" +] From 3eea0eb3fd62cdd6efc483f47341113f0765d833 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:11:46 +0530 Subject: [PATCH 10/29] shift the contact manager into the sensor manager --- .../managers/sensors/contact/config.py | 25 + .../sensors/contact/contact_manager.py | 551 ++++++++++++++++++ .../managers/sensors/contact/kernel.py | 90 +++ 3 files changed, 666 insertions(+) create mode 100644 genesis_forge/managers/sensors/contact/config.py create mode 100644 genesis_forge/managers/sensors/contact/contact_manager.py create mode 100644 genesis_forge/managers/sensors/contact/kernel.py 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] + ) From 31c7f7da6823ccdda03a7837a27d663e5a2603bf Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:18:28 +0530 Subject: [PATCH 11/29] add velocity action manager --- .../action/velocity_action_manager.py | 360 ++++++++++++++++++ 1 file changed, 360 insertions(+) create mode 100644 genesis_forge/managers/action/velocity_action_manager.py 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 From bf01c74e0bcd2624ad60a579b0defa7c9bd82ba8 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:18:47 +0530 Subject: [PATCH 12/29] add force action manager --- .../managers/action/force_action_manager.py | 360 ++++++++++++++++++ 1 file changed, 360 insertions(+) create mode 100644 genesis_forge/managers/action/force_action_manager.py 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 From 35dcd41d2fc3df7dcc9c8974564109625236bdbe Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:20:09 +0530 Subject: [PATCH 13/29] add hybrid action manager, to be used when a robot has both more than one type of control --- .../managers/action/hybrid_action_manager.py | 436 ++++++++++++++++++ 1 file changed, 436 insertions(+) create mode 100644 genesis_forge/managers/action/hybrid_action_manager.py 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 From b62a9cd6ca050a2235e1f0f5fba0304c31f01a5c Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:23:43 +0530 Subject: [PATCH 14/29] add ros2-env import and type hint --- genesis_forge/managers/command/command_manager.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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, ): From 05145c09cd677f8504bd2b3bc2d89fb6a1cd24bc Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:24:06 +0530 Subject: [PATCH 15/29] add force within limits action manager --- .../managers/action/force_within_limits.py | 102 ++++++++++++++++++ 1 file changed, 102 insertions(+) create mode 100644 genesis_forge/managers/action/force_within_limits.py 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 From 87e436ef96e920e56835a8af8ed882ba86a2e76a Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:25:34 +0530 Subject: [PATCH 16/29] add position command manager --- .../managers/command/position_command.py | 235 ++++++++++++++++++ 1 file changed, 235 insertions(+) create mode 100644 genesis_forge/managers/command/position_command.py 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}") + From 1cd1ae29b938e1f6d3246cbf519ad74dfcaad9b2 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:25:43 +0530 Subject: [PATCH 17/29] add pose command manager --- .../managers/command/pose_command.py | 246 ++++++++++++++++++ 1 file changed, 246 insertions(+) create mode 100644 genesis_forge/managers/command/pose_command.py 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}") From 51fecfa4a720b164f28b73b3d0105b2084637a42 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:27:44 +0530 Subject: [PATCH 18/29] add ros2-env import, modify type hint and call, shift any genesis-specific code to a conditional block --- .../managers/command/velocity_command.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) 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, From 442bde53959d51d1b5bf502dac7997d4ea9f8fee Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:32:05 +0530 Subject: [PATCH 19/29] add control type atttibute for each dof used, add properties for control type specific dof indices and dof names, add support for drone entities, add support for direct deploymenet to ros using the ros2-topic-based-control --- .../managers/actuator/actuator_manager.py | 344 +++++++++++++++++- 1 file changed, 332 insertions(+), 12 deletions(-) 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: From 42466272ee5bbc380ef4e8d43c2cb7722c4e063c Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:35:30 +0530 Subject: [PATCH 20/29] add sensor imports --- genesis_forge/managed_env.py | 54 +++++++++++++++++++++++++++++++++--- 1 file changed, 50 insertions(+), 4 deletions(-) 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, From 9b5963dc6e2fda53337e3a19e1271fcfeb79953e Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:38:05 +0530 Subject: [PATCH 21/29] fix all init to expose all submodules --- genesis_forge/__init__.py | 4 ++++ genesis_forge/managers/__init__.py | 20 ++++++++++++++++++-- genesis_forge/managers/action/__init__.py | 8 ++++++++ genesis_forge/managers/command/__init__.py | 4 ++++ 4 files changed, 34 insertions(+), 2 deletions(-) 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/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/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", ] From 4048edc393bfb604e9babf6791523adb5334d2ad Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:40:29 +0530 Subject: [PATCH 22/29] add ros2 publishers subscribers,and other ros2 specific componenets --- genesis_forge/managers/entity_manager.py | 47 +++++++++++++++++++++--- 1 file changed, 41 insertions(+), 6 deletions(-) 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) From 2f65b8bc27e7dc23fc4e4f7b76f9dac0a2208b97 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:41:26 +0530 Subject: [PATCH 23/29] add COM shift for domain randomisation --- genesis_forge/mdp/reset.py | 61 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) 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, + ) From 200ea481de556830ba1534ceb2c0e047b72fa03f Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:42:41 +0530 Subject: [PATCH 24/29] add new rewards for position-command tracking and pose-command tracking, add a rudementary imu-sensor based reward for rewarding stability --- genesis_forge/mdp/rewards.py | 272 ++++++++++++++++++++++++++++++++++- 1 file changed, 268 insertions(+), 4 deletions(-) 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 """ From d027162d89d19095dc2caaed41de3c2d643e55ca Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:45:29 +0530 Subject: [PATCH 25/29] add a simple example for using the skrl library along with the newly added sensor manager --- examples/skrl_with_sensors/README.md | 40 ++++ examples/skrl_with_sensors/environment.py | 280 ++++++++++++++++++++++ examples/skrl_with_sensors/eval.py | 87 +++++++ examples/skrl_with_sensors/train.py | 196 +++++++++++++++ 4 files changed, 603 insertions(+) create mode 100644 examples/skrl_with_sensors/README.md create mode 100644 examples/skrl_with_sensors/environment.py create mode 100644 examples/skrl_with_sensors/eval.py create mode 100644 examples/skrl_with_sensors/train.py 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() From b72c2768c66b99965480e73329e5ab8054ab4bec Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:50:04 +0530 Subject: [PATCH 26/29] modify observation aggregation to work with sensor data and proprioception, add a function to convert from observation dict to a gymnasium space, modify the action space to to have multiple gymnaisum box spaces one for each observation type --- genesis_forge/managers/observation_manager.py | 87 ++++++++++++++----- 1 file changed, 63 insertions(+), 24 deletions(-) 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") From 8de44378572ecc8043c816a1c677127a36464e25 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:51:22 +0530 Subject: [PATCH 27/29] add a ros2 manager env for direct deployment to ros --- genesis_forge/ros2_env.py | 292 ++++++++++++++++++++++ genesis_forge/ros2_managed_env.py | 401 ++++++++++++++++++++++++++++++ 2 files changed, 693 insertions(+) create mode 100644 genesis_forge/ros2_env.py create mode 100644 genesis_forge/ros2_managed_env.py 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() From e9b6a7be2e24bc9a23314827a7169762f9405cb4 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:52:00 +0530 Subject: [PATCH 28/29] cleanup --- genesis_forge/managers/action/position_action_manager.py | 1 + 1 file changed, 1 insertion(+) 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}") From b5291fb7c5a3b2a5608f0412a9867f02cd0d3146 Mon Sep 17 00:00:00 2001 From: "vybhav@gmail.com" Date: Sun, 23 Nov 2025 12:52:52 +0530 Subject: [PATCH 29/29] modify .gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) 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]