robot arm pushing task simulator

main.py

import gymnasium as gym
import panda_gym
from time import sleep

# see https://panda-gym.readthedocs.io/en/latest/ for available environments
env = gym.make("PandaPush-v3", render_mode="human")
observation, info = env.reset()

def controller(observation_vector, achieved_goal, desired_goal):
    pass

while True:

    # observation vector: end effector position (x,y,z), end effector velocity (x,y,z), object position (x,y,z), object rotation (x,y,z), object velocity (x,y,z), object angular velocity (x,y,z)
    observation_vector = observation["observation"][0:18]

    # achieved_goal: object position (x,y,z)
    achieved_goal = observation["achieved_goal"][0:3]

    # desired_goal: desired object position (x,y,z)
    desired_goal = observation["desired_goal"][0:3]

    # action: end effector position (x,y,z)
    action = controller(observation_vector, achieved_goal, desired_goal)

    observation, reward, terminated, truncated, info = env.step(action)

    if terminated or truncated:
        # observation, info = env.reset()
        pass
    
    sleep(0.05)

env.close()