Adds some new features

.gitignore

@@ -1,3 +1,4 @@
+.vscode/
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[codz]

examples/rough_terrain/environment.py

@@ -13,6 +13,7 @@
     VelocityCommandManager,
     TerrainManager,
     ContactManager,
+    ActuatorManager
 )
 from genesis_forge.mdp import reset, rewards, terminations
 

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

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)
+