Type hints, docstrings & refactor

[rainx0r]

• Type hinted all envs, utils and policies
• Added static type checking to pre-commit with MyPy
  • Added some type conversions to some env internals to ensure static type checking passes
• Added docstrings to all utils and SawyerXYZEnv
• Switched all unconverted mujoco calls to named access for readability and maintainability
• Moved assert_task_is_set decorator to be closed in SawyerXYZEnv to avoid a cyclical import
• Fixed NumPy DeprecationWarnings and updated all Box spaces to np.float64 to bypass casting and avoid warnings from latest Gymnasium
• Refactored env_dict.py to make it simpler to edit, and also deduplicated it (there were some duplicate keys)
• Made some changes in env XML files for future compatibility with MuJoCo 3. These changes do not affect behaviour in MuJoCo 2

Tested on Python 3.10.14

[Kallinteris-Andreas]

Regarding the changes to the XML model files, I would like to see verification using the check_environments_match() function, for multiple mujoco versions (100k steps per version should be enough)

[rainx0r]

Regarding the changes to the XML model files, I would like to see verification using the check_environments_match() function, for multiple mujoco versions (100k steps per version should be enough)

All right, perhaps I need to clarify my XML changes.

The main change I'm referring to is fixing a contype value typo (1s isn't a valid value, contype must be an integer) which is silently ignored in MuJoCo 2 it seems but hard checked in MuJoCo 3. This is present in shelf_dependencies.xml, which affects shelf-place-v2. I found this while trying to get Metaworld to work with MuJoCo 3.

There are two more XML changes which are adding a missing name annotation to two geoms that are important for stick-pull-v2/stick-push-v2 and soccer-v2 respectively. This is to fix a bug I found with static type checking. If you look at those envs currently (soccer, stick-push and stick-pull), they access the touching_object property. Except: there is no such property. If you look at SawyerXYZEnv, you can see that touching_main_object is the property, and touching_object is a function. Meaning, self.touching_object is an object, not None, so when converted to a boolean it's always True, while really what those envs wanted to do was call self.touching_main_object (which is what the rest of the envs do, e.g. push-v2). So, if we switch those envs with the bug ti self.touching_main_object, suddenly they fail, because if you look at its definition in SawyerXYZEnv, it essentially calls self.touching_object(id) on the id of the "main object" geom. Except that geom is not defined for those envs. This geom is defined with a name='objGeom' annotation in all other envs and that's the one SawyerXYZEnv tries to find in the MuJoCo model. So I added that annotation to the soccer ball and the stick to fix those envs. This means that those envs indeed did not have correct reward calculation before, from what I understand.

However, just to make sure, I also just made a quick verification script in a new branch on my fork of Metaworld here. I basically just added some new envs without my fixes to either touching_object -> touching_main_object in the env .py file, or the XML files. And then I used the recommended function to check, and on my machine at least on Python 3.10 and MuJoCo 2.3.7 it passes.

[reginald-mclean]

one issue with the type-hinting is that during the build tests, the python 3.8 build fails with message: TypeError: Type subscription requires python >= 3.9

[Kallinteris-Andreas]

I would like the xml_file changes moved in their own separate PR, and validate them for all supported mujoco versions (2.2.0 up to 2.3.7)

[reginald-mclean]

@Kallinteris-Andreas Why split them? The tests pass as expected. The only difference is now things will be calculated correctly for the affected environment reward functions.

[Kallinteris-Andreas]

In general it is best practice to have pull requests contain one type of change.
And I'm sensitive to the changes to the XML model files as they can have been reproducabily conseqences if wrong

[rainx0r]

In general it is best practice to have pull requests contain one type of change. And I'm sensitive to the changes to the XML model files as they can have been reproducabily conseqences if wrong

I totally understand that but:

1. One of the changes is fixing a value that is invalid in any version of MuJoCo, but is not validated / prevents the model from loading until 3.0
2. The other changes don't affect the model at all, they simply add name properties to two geoms so python code (that wasn't being called before because of a typo / bug) can reach them.

Sure, 1. could lead to a reproducibility issue because maybe an invalid value is treated as contype="0" whereas I set it to contype="1" since the invalid value was "1s" and I assume it's a typo for "1". But I did verify that the environments match, as you suggested with this script. And I could run it for more MuJoCo versions if you wish. Perhaps this change is unrelated to the rest of the PR, but it's a minor typo fix, so does it really deserve its own PR?

As for 2., those changes are absolutely relevant to this PR. I found that bug through static type analysis. If I revert the XML changes to remove those name properties, the tests now fail, because the envs are actually looking for objects with that name. So then I would have to also reintroduce the bug (self.touching_main_object -> self.touching_object) to the envs to prevent them from trying to (rightfully) look for objects with that name. Except, then the MyPy checks would fail. And then I would have to remove MyPy checks from the pre-commit, which imo would almost defeat the point of type annotations in a way. Their benefit is the fact that we can do static type analysis and find bugs like these. This change also has nothing to do with MuJoCo, it just so happens that to make the now-reachable python code work I needed to add a name property to a couple geoms in the XML--something that does not affect model behaviour.

[Kallinteris-Andreas]

@evangelos-ch at least validate the XML file changes for all supported mujoco versions as described in a previous comment

[reginald-mclean]

@Kallinteris-Andreas Mujoco 2.2.0-2.2.2 throw errors unrelated to this PR, we will test for 2.3.0->2.3.7

[rainx0r]

As requested, I added a higher-level script in my xml-changes-verification branch that runs the checking script for all MuJoCo versions between 2.2.0 and 2.3.7.

The script passes for all 2.3.X versions with info comparison disabled. With info comparison enabled, it expectedly fails at 5000 steps because grasp_success now correctly turns to 1.0 since it can actually correctly process whether the arm is touching the object, while on the unmodified environments with this bug it doesn't.

The test fails for 2.2.0 / 2.2.1 / 2.2.2, but not even at the env comparison step. It seems as though Metaworld isn't compatible with those MuJoCo versions, which lines up with my own understanding and experience maintaining the benchmark. When we first worked to migrate off mujoco-py, MuJoCo 2.3.4 was already out.

Also for reference, these are the scripts used for verification:

• Python script using the suggested Gymnasium utility
• Higher-level shell script that calls it for each MuJoCo version

I'm uploading logs from both runs (with info comparison and without).

• xml_verification.log
• xml_verification_with_info.log

[reginald-mclean]

@Kallinteris-Andreas ^

[Kallinteris-Andreas]

LGTM with regards to the XML file changes.