Unstable behavior of box in Ant environment #497

wladekpal · 2024-07-01T20:37:04Z

wladekpal
Jul 1, 2024

Hi, I've been trying to replicate one of the environments from Data-Efficient Hierarchical Reinforcement Learning in brax, specifically Ant Push, which is an Ant env with some walls and a movable block, that has to be pushed out of the way to reach the goal. I've created the following xml file:

<mujoco model="ant">
	<compiler angle="degree" coordinate="local" inertiafromgeom="true" />
	<option timestep="0.01" iterations="4"/>
	<custom>
		<numeric data="0.0 0.0 0.75 1.0 0.0 0.0 0.0 0.0 1.0 0.0 -1.0 0.0 -1.0 0.0 1.0 0.0 0.0 0.0 0.0 0.0" name="init_qpos" />
		<numeric data="1000" name="constraint_limit_stiffness" />
		<numeric data="4000" name="constraint_stiffness" />
		<numeric data="10" name="constraint_ang_damping" />
		<numeric data="0.5" name="joint_scale_pos" />
		<numeric data="0.2" name="joint_scale_ang" />
		<numeric data="0.0" name="ang_damping" />
		<numeric data="1" name="spring_mass_scale" />
		<numeric data="1" name="spring_inertia_scale" />
		<numeric data="15" name="solver_maxls" />
	</custom>
	<default>
		<joint armature="1" damping="1" limited="true" />
		<geom contype="0" conaffinity="0" condim="3" density="5.0" friction="1 0.5 0.5" solimp=".995 .995 .01" />
	</default>
	<asset>
		<texture builtin="gradient" height="100" rgb1="1 1 1" rgb2="0 0 0" type="skybox" width="100" />
		<texture builtin="flat" height="1278" mark="cross" markrgb="1 1 1" name="texgeom" random="0.01" rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" type="cube" width="127" />
		<texture builtin="checker" height="100" name="texplane" rgb1="0 0 0" rgb2="0.8 0.8 0.8" type="2d" width="100" />
		<material name="MatPlane" reflectance="0.5" shininess="1" specular="1" texrepeat="60 60" texture="texplane" />
		<material name="geom" texture="texgeom" texuniform="true" />
	</asset>
	<worldbody>
		<light cutoff="100" diffuse="1 1 1" dir="-0 0 -1.3" directional="true" exponent="1" pos="0 0 1.3" specular=".1 .1 .1" />
		<geom conaffinity="1" condim="3" material="MatPlane" name="floor" pos="0 0 0" size="100 100 40" type="plane" />
		<body name="torso" pos="0 0 0.75">
			<camera name="track" mode="trackcom" pos="0 -3 0.3" xyaxes="1 0 0 0 0 1" />
			<geom name="torso_geom" pos="0 0 0" size="0.25" type="sphere" contype="1" conaffinity="1"/>
			<joint armature="0" damping="0" limited="false" margin="0.01" name="root" pos="0 0 0" type="free" />
			<body name="front_left_leg" pos="0 0 0">
				<geom fromto="0.0 0.0 0.0 0.2 0.2 0.0" name="aux_1_geom" size="0.08" type="capsule" contype="1" conaffinity="1" />
				<body name="aux_1" pos="0.2 0.2 0">
					<joint axis="0 0 1" name="hip_1" pos="0.0 0.0 0.0" range="-30 30" type="hinge" />
					<geom fromto="0.0 0.0 0.0 0.2 0.2 0.0" name="left_leg_geom" size="0.08" type="capsule" contype="1" conaffinity="1"/>
					<body pos="0.2 0.2 0">
						<joint axis="-1 1 0" name="ankle_1" pos="0.0 0.0 0.0" range="30 70" type="hinge" />
						<geom fromto="0.0 0.0 0.0 0.4 0.4 0.0" name="left_ankle_geom" size="0.08" type="capsule" contype="1" conaffinity="1"/>
						<geom name="left_foot_geom" pos="0.4 0.4 0" size="0.08" type="sphere" mass="0" contype="1" conaffinity="1"/>
					</body>
				</body>
			</body>
			<body name="front_right_leg" pos="0 0 0">
				<geom fromto="0.0 0.0 0.0 -0.2 0.2 0.0" name="aux_2_geom" size="0.08" type="capsule" contype="1" conaffinity="1" />
				<body name="aux_2" pos="-0.2 0.2 0">
					<joint axis="0 0 1" name="hip_2" pos="0.0 0.0 0.0" range="-30 30" type="hinge" />
					<geom fromto="0.0 0.0 0.0 -0.2 0.2 0.0" name="right_leg_geom" size="0.08" type="capsule" contype="1" conaffinity="1"/>
					<body pos="-0.2 0.2 0">
						<joint axis="1 1 0" name="ankle_2" pos="0.0 0.0 0.0" range="-70 -30" type="hinge" />
						<geom fromto="0.0 0.0 0.0 -0.4 0.4 0.0" name="right_ankle_geom" size="0.08" type="capsule" contype="1" conaffinity="1"/>
						<geom name="right_foot_geom" pos="-0.4 0.4 0" size="0.08" type="sphere" mass="0" contype="1" conaffinity="1"/>
					</body>
				</body>
			</body>
			<body name="back_leg" pos="0 0 0">
				<geom fromto="0.0 0.0 0.0 -0.2 -0.2 0.0" name="aux_3_geom" size="0.08" type="capsule" contype="1" conaffinity="1"/>
				<body name="aux_3" pos="-0.2 -0.2 0">
					<joint axis="0 0 1" name="hip_3" pos="0.0 0.0 0.0" range="-30 30" type="hinge" />
					<geom fromto="0.0 0.0 0.0 -0.2 -0.2 0.0" name="back_leg_geom" size="0.08" type="capsule" contype="1" conaffinity="1"/>
					<body pos="-0.2 -0.2 0">
						<joint axis="-1 1 0" name="ankle_3" pos="0.0 0.0 0.0" range="-70 -30" type="hinge" />
						<geom fromto="0.0 0.0 0.0 -0.4 -0.4 0.0" name="third_ankle_geom" size="0.08" type="capsule" contype="1" conaffinity="1"/>
						<geom name="third_foot_geom" pos="-0.4 -0.4 0" size="0.08" type="sphere" mass="0" contype="1" conaffinity="1"/>
					</body>
				</body>
			</body>
			<body name="right_back_leg" pos="0 0 0">
				<geom fromto="0.0 0.0 0.0 0.2 -0.2 0.0" name="aux_4_geom" size="0.08" type="capsule" contype="1" conaffinity="1"/>
				<body name="aux_4" pos="0.2 -0.2 0">
					<joint axis="0 0 1" name="hip_4" pos="0.0 0.0 0.0" range="-30 30" type="hinge" />
					<geom fromto="0.0 0.0 0.0 0.2 -0.2 0.0" name="rightback_leg_geom" size="0.08" type="capsule" contype="1" conaffinity="1"/>
					<body pos="0.2 -0.2 0">
						<joint axis="1 1 0" name="ankle_4" pos="0.0 0.0 0.0" range="30 70" type="hinge" />
						<geom fromto="0.0 0.0 0.0 0.4 -0.4 0.0" name="fourth_ankle_geom" size="0.08" type="capsule" contype="1" conaffinity="1"/>
						<geom name="fourth_foot_geom" pos="0.4 -0.4 0" size="0.08" type="sphere" mass="0" contype="1" conaffinity="1"/>
					</body>
				</body>
			</body>
		</body>
		
		<geom name="block_0_0" pos="-8.000000 -8.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_0_1" pos="0.000000 -8.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_0_2" pos="8.000000 -8.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_0_3" pos="16.000000 -8.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_0_4" pos="24.000000 -8.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_1_0" pos="-8.000000 0.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_1_3" pos="16.000000 0.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_1_4" pos="24.000000 0.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_2_0" pos="-8.000000 8.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<body name="movable_2_2" pos="8.000000 8.000000 2.100000">
			<geom name="block_2_2" pos="0 0 0" size="3.600000 3.600000 2.00000" type="box" material="" mass="1" contype="1" conaffinity="1" rgba="0.9 0.1 0.1 1" />
			<joint armature="0" axis="1 0 0" damping="0.0" limited="false" range="-32.000000 32.000000" margin="0.01" name="movable_x_2_2" pos="0 0 0" type="slide" />
			<joint armature="0" axis="0 1 0" damping="0.0" limited="false" range="-32.000000 32.000000" margin="0.01" name="movable_y_2_2" pos="0 0 0" type="slide" />
			<joint armature="0" axis="0 0 1" damping="0.0" limited="false" range="-32.000000 32.000000" margin="0.01" name="movable_z_2_2" pos="0 0 0" type="slide" />
		</body>
		<geom name="block_2_4" pos="24.000000 8.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_3_0" pos="-8.000000 16.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_3_1" pos="0.000000 16.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_3_3" pos="16.000000 16.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_3_4" pos="24.000000 16.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_4_0" pos="-8.000000 24.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_4_1" pos="0.000000 24.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_4_2" pos="8.000000 24.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_4_3" pos="16.000000 24.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />
		<geom name="block_4_4" pos="24.000000 24.000000 2.000000" size="4.000000 4.000000 2.000000" type="box" material="" contype="1" conaffinity="1" rgba="0.4 0.4 0.4 1" />

    <body name="target" pos="0 0 0.01">
		<joint armature="0" axis="1 0 0" damping="200" limited="true" name="target_x" pos="0 0 0" range="-20 20" stiffness="0" type="slide" />
		<joint armature="0" axis="0 1 0" damping="200" limited="true" name="target_y" pos="0 0 0" range="-20 20" stiffness="0" type="slide" />
		<geom conaffinity="0" contype="0" name="target" pos="0 0 0" size=".9" type="sphere" mass="inf" />
	</body>
	</worldbody>
	<actuator>
		<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_4" gear="150" />
		<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_4" gear="150" />
		<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_1" gear="150" />
		<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_1" gear="150" />
		<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_2" gear="150" />
		<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_2" gear="150" />
		<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_3" gear="150" />
		<motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_3" gear="150" />
	</actuator>
</mujoco>

This is basically exactly the same as ant.xml in brax, but with addition of walls and movable block. The problem I encountered is that in the spring backend the block doesn't sit on the plane like it's supposed to but instead it constantly jumps and wobbles. This is very weird for me, especially considering that interaction between block and ant model seem to be ok. Here is a video of this behavior:
https://github.com/google/brax/assets/56890546/87861a18-fd23-4ab8-bcf8-f940c00626e3

After some tinkering I've noticed that removing <numeric data="1" name="spring_inertia_scale" /> from custom tag in the xml file makes the movable block behave in a desired way, but instead it makes the ant model break down completely as can be seen in the this video:
https://github.com/google/brax/assets/56890546/850b8dd7-786b-4197-ab97-16e5618e8b46

I've tested it for positional pipeline, and the behavior is somewhat similar, but the block is ejected into air, instead of wobbling (and removing spring_inertia_scale has the exact same effect).

My understanding is that <numeric data="1" name="spring_inertia_scale" /> is added only to the brax version of ant, as I haven't seen similiar thing in the original mujoco implementation of ant, but I'm not really sure what is the reason it is added and why Ant brakes down without it.

I would appreciate some help with this issue, as I am not very experienced in mujoco and brax, so it is quite hard for me to understand really what is the problem here.

Answered by erikfrey

Jul 2, 2024

I'd recommend trying this with the MJX pipeline - the spring pipeline is difficult to tune, especially for large inertia ratios (small ant vs big block).

View full answer

erikfrey · 2024-07-02T19:32:40Z

erikfrey
Jul 2, 2024
Maintainer

I'd recommend trying this with the MJX pipeline - the spring pipeline is difficult to tune, especially for large inertia ratios (small ant vs big block).

4 replies

wladekpal Jul 5, 2024
Author

Thanks for quick response, @erikfrey, I tested it and with MJX pipeline the environment seems stable. I have one more question though.
I've trained a model on the same ant environment with both spring and mjx pipelines (using default ant.xml taken from brax). While it did learn well on spring it completely failed on MJX. So my question is, whether MJX has been observed to be significantly harder than other pipelines, and if not, should I somehow tune the default ant.xml from brax repo to be suitable for MJX ?
It is my understanding that MJX aims to recreate default mujoco behavior and from what I could gather .xml files for the same models (like ant) differ between mujoco and brax, which leads me to belive that my problem may be in using the same .xml on both spring and MJX pipelines.

erikfrey Jul 10, 2024
Maintainer

The MJX pipeline models contact restitution and friction in a more detailed way that probably yields more diverse dynamics - so yes it's probably a bit harder of an exploration problem, but not by much. You should be able to get a useful policy in about the same number of training steps as before, but you may need different hyperparamters.

You should be able to use the provided ant.xml from the brax repo as is. I don't have a set of known good PPO hyperparams for ant handy, but I suspect you could make progress starting from the humanoid hparams.

erikfrey Jul 10, 2024
Maintainer

Oh you know what, I take that back - you may want to tweak the XML for MJX. I would try changing the <option> in ant.xml tag to look like this:

  <option timestep="0.01" iterations="1" ls_iterations="4">
    <flag eulerdamp="disable"/>
  </option>

This means you use the Newton solver, with only 1 iteration and 4 linesearch iterations. And disable implicit integration with respect to joint damping - should make training times faster. Try these settings but confirm the simulation looks stable - you may need to nudge iterations and ls_iterations upwards a bit if it's not stable.

wladekpal Jul 11, 2024
Author

Ok, I've changed iterations to 4, and this led to much better behavior. Additionally I've increased the density of all of the model componenets and these 2 changes seemed to get me where I wanted. Thank you for your help @erikfrey.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Unstable behavior of box in Ant environment #497

{{title}}

Replies: 1 comment 4 replies

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

Select a reply

Unstable behavior of box in Ant environment #497

wladekpal Jul 1, 2024

Replies: 1 comment · 4 replies

erikfrey Jul 2, 2024 Maintainer

wladekpal Jul 5, 2024 Author

erikfrey Jul 10, 2024 Maintainer

erikfrey Jul 10, 2024 Maintainer

wladekpal Jul 11, 2024 Author

wladekpal
Jul 1, 2024

Replies: 1 comment 4 replies

erikfrey
Jul 2, 2024
Maintainer

wladekpal Jul 5, 2024
Author

erikfrey Jul 10, 2024
Maintainer

erikfrey Jul 10, 2024
Maintainer

wladekpal Jul 11, 2024
Author