running_env.py

import math

import numpy as np
import gymnasium as gym
from gymnasium import spaces

class RunningEnv(gym.Env):
    """The environment for training the model to run."""
    metadata = {"render_modes": ["human"]}

    def __init__(self, ragdoll):
        """Initialize the environment."""
        self.ragdoll = ragdoll
        self.space = self.ragdoll.space
        self.max_steps = 500
        self.current_step = 0

        self.previous_x = 0

        # initializes the number of inputs to an action
        num_joints = len(self.ragdoll.motors)
        self.action_space = spaces.Box(
            low=-1.0, high=1.0, shape=(num_joints,), dtype=np.float32
        )

        # number of inputs to a state
        obs_dim = len(self.ragdoll.parts) * 6
        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(obs_dim,), dtype=np.float32
        )

    def reset(self, seed=None, options=None):
        """Reset the environment."""
        super().reset(seed=seed)
        self.current_step = 0
        self.ragdoll.reset()
        return self.get_state(), {}

    def get_state(self):
        """
        Gets the current state of the ragdoll to pass in to the model.
        :return: a numpy array with 6 values for each body part (values shown below)
        """
        state = []
        for part in self.ragdoll.parts.values():
            state.extend([part.position.x, part.position.y,
                          part.velocity.x, part.velocity.y,
                          part.angle, part.angular_velocity])
        return np.array(state, dtype=np.float32)

    def step(self, action):
        """
        applies actions, computes rewards, and steps physics simulation for one step
        :param action: the action to apply to the ragdoll
        :return: obs, reward, terminated, truncated, info
        """
        self.ragdoll.apply_action(action)
        dt = 1 / 60.0
        self.space.step(dt)
        self.current_step += 1

        reward = self.compute_reward()
        terminated = self.is_fallen()
        truncated = self.current_step >= self.max_steps

        obs = self.get_state()
        info = {}

        return obs, reward, terminated, truncated, info

    def compute_reward(self):
        """Calculates the reward based on the state of the ragdoll."""
        torso = self.ragdoll.parts["torso"]

        forward_velocity = torso.velocity.x
        velocity_reward = forward_velocity * 0.1  # reward for forward velocity

        """tilt_angle = abs(torso.angle)
        upright_reward = max(0, 1.0 - (tilt_angle / 0.5) ** 2)  # exponential penalty for spinning

        if tilt_angle > 1.0:
            upright_reward = -5.0"""

        x_progress = (torso.position.x - self.previous_x) * 0.5  # reward for moving forward
        self.previous_x = torso.position.x
        angular_velocity_penalty = -10 * abs(torso.angular_velocity)  # penalty for falling

        total_reward = velocity_reward + x_progress + angular_velocity_penalty

        if self.is_fallen():
            # penalty for falling off the side
            total_reward -= 10

        return total_reward

    def is_fallen(self):
        """Determines if the ragdoll has fallen out of the environment."""
        torso = self.ragdoll.parts["torso"]

        return torso.position.y > 800 or torso.position.y < 0

    def render(self):
        """Renders the ragdoll."""
        self.ragdoll.render()
        self.ragdoll.clock.tick(60)