import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import gymnasium as gym

# Make Pendulum environment
env = gym.make('Pendulum-v1', g=9.81)  # g is gravitational acceleration, default 10
state_dim = env.observation_space.shape[0]   # dimension of state (should be 3 for Pendulum)
action_dim = env.action_space.shape[0]       # dimension of action (1 for Pendulum)
max_action = float(env.action_space.high[0]) # max torque (=2.0)