Module pearl.user_envs.wrappers.gym_avg_torque_cost

Expand source code
# pyre-ignore-all-errors

try:
    import gymnasium as gym
except ModuleNotFoundError:
    print("gymnasium module not found.")


class GymAvgTorqueWrapper(gym.Wrapper):
    r"""Sparse Reward wrapper for the Pendulum environment."""

    def __init__(self, env):
        super(GymAvgTorqueWrapper, self).__init__(env)

    def step(self, action):
        obs, reward, done, truncated, info = self.env.step(action)
        # assumes action is tensor
        cost = (action**2).mean()
        info["cost"] = cost
        return obs, reward, done, truncated, info

Classes

class GymAvgTorqueWrapper (env)

Sparse Reward wrapper for the Pendulum environment.

Wraps an environment to allow a modular transformation of the :meth:step and :meth:reset methods.

Args

env
The environment to wrap
Expand source code
class GymAvgTorqueWrapper(gym.Wrapper):
    r"""Sparse Reward wrapper for the Pendulum environment."""

    def __init__(self, env):
        super(GymAvgTorqueWrapper, self).__init__(env)

    def step(self, action):
        obs, reward, done, truncated, info = self.env.step(action)
        # assumes action is tensor
        cost = (action**2).mean()
        info["cost"] = cost
        return obs, reward, done, truncated, info

Ancestors

  • gymnasium.core.Wrapper
  • gymnasium.core.Env
  • typing.Generic

Methods

def step(self, action)

Uses the :meth:step of the :attr:env that can be overwritten to change the returned data.

Expand source code
def step(self, action):
    obs, reward, done, truncated, info = self.env.step(action)
    # assumes action is tensor
    cost = (action**2).mean()
    info["cost"] = cost
    return obs, reward, done, truncated, info