Feature/sensor observation, imu reward functions

examples/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/
+```

examples/sensors/environment.py

@@ -0,0 +1,299 @@
+"""
+Simplified Go2 Locomotion Environment using managers to handle everything.
+"""
+
+import torch
+import genesis as gs
+
+from genesis_forge import ManagedEnvironment
+from genesis_forge import EnvMode
+from genesis_forge.managers import (
+    RewardManager,
+    TerminationManager,
+    EntityManager,
+    ObservationManager,
+    ActuatorManager,
+    PositionActionManager,
+)
+from genesis_forge.mdp import reset, rewards, terminations, observations
+
+
+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,
+    ):
+        super().__init__(
+            num_envs=num_envs,
+            dt=dt,
+            max_episode_length_sec=max_episode_length_s,
+            max_episode_random_scaling=0.1,
+        )
+
+        # 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
+        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,
+            ),
+        )
+        self.imu = self.scene.add_sensor(
+            gs.sensors.IMU(
+                entity_idx=self.robot.idx,
+                link_idx_local=self.robot.links[0].idx_local,
+                pos_offset=(0.0, 0.0, 0.15),
+                # noise parameters
+                acc_cross_axis_coupling=(0.0, 0.01, 0.02),
+                gyro_cross_axis_coupling=(0.03, 0.04, 0.05),
+                acc_noise=(0.01, 0.01, 0.01),
+                gyro_noise=(0.01, 0.01, 0.01),
+                acc_random_walk=(0.001, 0.001, 0.001),
+                gyro_random_walk=(0.001, 0.001, 0.001),
+                # delay=0.01,
+                jitter=0.01,
+                interpolate=True,
+                # visualize
+                draw_debug=True,
+            )
+        )
+
+        # Camera, for headless video recording
+        self.camera = self.scene.add_camera(
+            pos=(-2.5, -1.5, 1.0),
+            lookat=(0.0, 0.0, 0.0),
+            res=(1280, 720),
+            fov=40,
+            env_idx=0,
+            debug=True,
+        )
+
+    def config(self):
+        """
+        Configure the environment managers
+        """
+        ##
+        # Robot manager
+        # i.e. what to do with the robot when it is reset
+        self.robot_manager = EntityManager(
+            self,
+            entity_attr="robot",
+            on_reset={
+                # Reset the robot's initial position
+                "position": {
+                    "fn": reset.position,
+                    "params": {
+                        "position": INITIAL_BODY_POSITION,
+                        "quat": INITIAL_QUAT,
+                        "zero_velocity": True,
+                    },
+                },
+            },
+        )
+
+        ##
+        # Joint Actions
+        self.actuator_manager = ActuatorManager(
+            self,
+            joint_names=[
+                "FL_.*_joint",
+                "FR_.*_joint",
+                "RL_.*_joint",
+                "RR_.*_joint",
+            ],
+            default_pos={
+                ".*_hip_joint": 0.0,
+                "FL_thigh_joint": 0.8,
+                "FR_thigh_joint": 0.8,
+                "RL_thigh_joint": 1.0,
+                "RR_thigh_joint": 1.0,
+                ".*_calf_joint": -1.5,
+            },
+            kp=20,
+            kv=0.5,
+        )
+        self.action_manager = PositionActionManager(
+            self,
+            scale=0.25,
+            clip=(-100.0, 100.0),
+            use_default_offset=True,
+            actuator_manager=self.actuator_manager,
+        )
+
+        ##
+        # Observations
+        self.observation_manager = ObservationManager(
+            self,
+            history_len=2,
+            cfg={
+                "angle_velocity": {
+                    "fn": lambda env: self.robot_manager.get_angular_velocity(),
+                    "scale": 0.25,
+                },
+                "linear_velocity": {
+                    "fn": lambda env: self.robot_manager.get_linear_velocity(),
+                    "scale": 2.0,
+                },
+                "projected_gravity": {
+                    "fn": lambda env: self.robot_manager.get_projected_gravity(),
+                },
+                "dof_position": {
+                    "fn": lambda env: self.action_manager.get_dofs_position(),
+                },
+                "dof_velocity": {
+                    "fn": lambda env: self.action_manager.get_dofs_velocity(),
+                    "scale": 0.05,
+                },
+                "imu": {
+                    "fn": observations.SensorObservation,
+                    "params": {
+                        "read": lambda: torch.cat(
+                            [self.imu.read().lin_acc, self.imu.read().ang_vel], dim=-1
+                        ),
+                        "frequency": 25,
+                    },
+                },
+                "actions": {
+                    "fn": lambda env: self.action_manager.get_actions(),
+                },
+            },
+        )
+
+        ##
+        # 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,
+                    },
+                },
+                "lin_acc_jitter": {
+                    "weight": -0.1,
+                    "fn": rewards.imu_lin_acc_jitter,
+                    "params": {
+                        "observation_manager": self.observation_manager,
+                        "obs_item_key": "imu",
+                        "ignore_gravity": True,
+                    },
+                },
+                "ang_vel_jitter": {
+                    "weight": -0.1,
+                    "fn": rewards.imu_ang_vel_jitter,
+                    "params": {
+                        "observation_manager": self.observation_manager,
+                        "obs_item_key": "imu",
+                    },
+                },
+            },
+        )
+
+        ##
+        # 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,
+                    },
+                },
+            },
+        )
+
+    def build(self):
+        super().build()
+        self.camera.follow_entity(self.robot)