Add keypoints mode, and fix order of operations in add_tee

[alexander-soare]

This PR does two mostly disconnected things:

1. Adds an observation mode where keypoints and agent position are returned, as an alternative to pixels_agent_pos which returns an image and the agent position. This sort of keypoints-only environment allows us to circumvent using image encoders in the policy, which makes for a more agile experimentation setting.

2. Fixes the order of angle and position setting in add_tee.

You visualize the keypoints with this script:

import cv2
import gymnasium as gym
import numpy as np

import gym_pusht

env = gym.make(
    "gym_pusht/PushT-v0",
    render_mode="rgb_array",
    obs_type="environment_state_agent_pos",
)
observation, info = env.reset()

for _ in range(1000):
    action = env.action_space.sample()
    observation, reward, terminated, truncated, info = env.step(action)
    image = env.render()
    if terminated or truncated:
        observation, info = env.reset()

    for kp in observation["environment_state"].reshape(8, 2):
        kp = np.round(kp / 512 * 680).astype(int)
        cv2.circle(image, tuple(kp), radius=4, color=(0, 0, 255), thickness=-1)
    cv2.imshow("window", image)
    k = cv2.waitKey(100)
    if k == ord("q"):
        break

env.close()

The reason 2 (fix order of ops in add_tee) is here, is because when generating the keypoints-only dataset using this script I found out that the keypoints weren't aligned with the tee. In fact, I also discovered that the reward calculation was incorrect, for the same reason. By switching the order of ops, I was able to rectify both issues.

[alexander-soare]

Although I have empirically determined that this is correct, I'd much rather understand it and add a comment explaining what's going on before merging this PR.

[qgallouedec]

If I remember correctly, the rotation is made around a point that is different from the one used for the position. Therefore, rotating the object changes its position. Having the code in this order makes more sense.

[alexander-soare]

Okay. I did some testing and it seemed that setting the position moves the bodies in the object frame, whereas setting the angle rotates the object frame in the world frame (consistent with you explanation). But, it was really weird... because I couldn't always get it to work that way.
I'm not sure I have the appetite to fully resolve this right now though... but if someone else has problems in the future, they will have to figure it out.

[qgallouedec]

is get_keypoints meant to be accessed outside the env?
Starting from gymnasium 1.0, you'll need to set env.unwrapped.my_func if you want to access a custom method that is part of your env. See https://github.com/Farama-Foundation/Gymnasium/blob/81b87efb9f011e975f3b646bab6b7871c522e15e/gymnasium/core.py#L312
It probably won't affect this PR but good to keep in mind.

[qgallouedec]

I'd also add a test here: https://github.com/alexander-soare/gym-pusht/blob/9d19f907f2c63328df75850920982be1b0cf7e95/tests/test_env.py#L13

[Cadene]

LGTM

[Cadene]

what an artist 🧑‍🎨

[alexander-soare]

is get_keypoints meant to be accessed outside the env? Starting from gymnasium 1.0, you'll need to set env.unwrapped.my_func if you want to access a custom method that is part of your env. See https://github.com/Farama-Foundation/Gymnasium/blob/81b87efb9f011e975f3b646bab6b7871c522e15e/gymnasium/core.py#L312 It probably won't affect this PR but good to keep in mind.

@qgallouedec Thanks. Yeah it is accessed in our dataset conversion script where we have to reconstruct the environment in order to access reward and now the keypoints. Since some prior methods are exposed and accessed for the same reason, I'm happy to leave it out of this PR if you are.

I'd also add a test here: https://github.com/alexander-soare/gym-pusht/blob/9d19f907f2c63328df75850920982be1b0cf7e95/tests/test_env.py#L13

Done.

[qgallouedec]

Sure! Gym v1 is alpha anyway, we still have plenty of time to migrate