Transport Environment

The transport environment extends legged_gym's LeggedRobot. The robot carries a payload that is pulled by a simulated partner; the high-level policy has to infer this pull from its arm joints and move accordingly. The same environment serves the teacher and the student; only the actor observation differs.

Robot and payload

We use ANYmal D with a 6-DoF DynaArm. Any quadruped with a 6-DoF arm should work.
The payload is a rigid box with two handles, attached to the end-effector by a fixed joint. Keep it a separate rigid body with dont_collapse="true", so you can randomize its mass and apply the wrench to it.
The arm's default posture holds the payload in front of the base. Arm joints are position controlled with PD gains K_p = 50, K_d = 4. Self-collisions are enabled.

Simulation settings

Physics time step	0.005 s
Decimation	4, so the policies run at 50 Hz (control dt = 0.02 s)
Episode length	10 s (500 control steps)
Terrain	Rough terrain from legged_gym's generator (slopes, stairs and discrete obstacles), with the terrain curriculum
Environments	4096

On rough terrain the LL policy needs its height scan. The terrain curriculum should measure progress against the wrench reference instead of legged_gym's velocity command.

Action space

The high-level action a ∈ ℝ⁹ has two parts:

a_0:3: base command (v_x, v_y, ω_z) for the frozen LL policy, after the limits and deadband. Not scaled.
a_3:9: arm joint targets q^arm_des = q^arm₀ + 0.5 · a_3:9.

At every control step, the processed base command goes to the LL policy, which returns the leg joint targets. Leg and arm targets are held over the 4 physics steps and tracked by joint PD controllers. The interaction wrench must be applied again before every physics step, because Isaac Gym clears applied forces after each simulation step.

Interaction wrench generation

The partner is simulated by a planar wrench on the payload. At every reset we sample a normalized target c̄ = (c̄_x, c̄_y, c̄_z) uniformly from [−1, 1]³; with probability 0.02 the episode gets zero wrench instead, so the robot also learns to stand still. Over the episode the wrench follows a piecewise-linear profile s(t): 10% ramp-up, 80% hold, 10% ramp-down.

c(t) = s(t) c̄, F_xy = B_F c_xy(t), Θ_z = B_T c_z(t), B_F = 40 N, B_T = 10 N·m

The wrench is defined in the base frame, rotated to the world frame and applied to the payload body at every physics step. The normalized c(t) plays three roles: it is the teacher's privileged observation, the intent estimator's regression target, and the velocity reference in the tracking rewards (40 N corresponds to 1 m/s, 10 N·m to 1 rad/s).

Domain randomization is also applied to the generated wrench. In our setup, a small periodic disturbance acts on the end-effector to mimic the partner's hand motion: a vertical force 2 sin(2πt) N and torques sin(2πt + π/4) and sin(2πt + π/2) N·m about the world x and y axes, with t the time since the episode started.

Observations

The student observes only arm proprioception over the last H control steps (Eq. (1) of the paper), stacked from oldest to newest, flattened, and zero-filled at reset. a^arm is the previous arm joint position command. The teacher additionally observes the current interaction wrench (Eq. (2)). Teacher and student both use an asymmetric actor-critic, whose critic additionally observes end-effector and payload position, velocity and acceleration. No observation normalization is used.

Group	Contents	Size (H = 8)
Student actor o^trans	q^arm (noise ±0.01), q̇^arm (noise ±1.5) and the previous arm command a^arm, over H steps	(6 + 6 + 6) × 8 = 144
Teacher actor o^T	o^trans plus the current wrench o^wrench, normalized as c(t)	144 + 3 = 147
Critic (both stages)	Teacher observation plus end-effector and payload position, velocity and acceleration	—

The student actor additionally receives the estimated wrench ŵ (3 values) from the intent estimator, so its network input is 144 + 3 = 147. That concatenation happens in the algorithm, not in the environment (see Training). Arm joint positions are measured relative to the default posture, and the noise values are bounds of uniform noise.

In the teacher task, the teacher observation is the actor input. In the student task, the student observation is the actor input, and the teacher observation is passed alongside it so the algorithm can query the teacher and read the true wrench.

Rewards

The rewards align the base motion with the wrench on the payload, like an admittance controller, while keeping the arm still and the payload level. The tracking terms are

r_force = exp(−‖F^xy_force / B_force − v^xy_base‖₂² / σ_force), r_torque = exp(−(Θ^yaw_torque / B_torque − ω^yaw_base)² / σ_torque)

and the payload height term is

r_h = (max(0, h_min − h^p_base) + max(0, h^p_base − h_max))²

with σ_force = σ_torque = 0.25, B_force = 40, B_torque = 10, h_min = 0.05 m and h_max = 0.25 m. v^xy_base and ω^yaw_base are the base linear velocity and yaw rate, and h^p_base is the payload height, all in the base frame. Since F/B_force = c_xy and Θ/B_torque = c_z, the tracking terms compare the base velocity directly with the normalized wrench.

The weights are those of Table I in the paper. Each weight is multiplied by the control interval dt = 0.02 s, which is what legged_gym does with the reward scales in the config.

Reward term	Expression	Weight (×dt)
Force tracking	r_force	−2.0 × 10⁻²
Torque tracking	r_torque	−1.0 × 10⁻²
Payload height	r_h	−3.0 × 10⁻³
Joint torque penalty	Σ_i τ_i²	−2.0 × 10⁻⁹
Joint DoF penalty	Σ_i q̇_i² + 0.00025 q̈_i²	−2.0 × 10⁻⁴
Action penalty	Σ_i (a_t,i − a_t−1,i)² + 0.5 a_t,i²	−2.0 × 10⁻⁴
Termination	𝕀[terminated]	−50.0 / dt
Arm posture penalty	‖q_arm − q_arm^default‖₂	−3.0 × 10⁻³

The tracking terms use the wrench reference instead of legged_gym's velocity commands. The payload height and the critic's end-effector terms need rigid-body states, which legged_gym does not acquire by default.

Terminations

Base contact with the ground.
Time-out after 10 s. Time-outs are bootstrapped.

Domain randomization

Payload mass	Uniform 0-10 kg, per environment
Base added mass	Uniform −2 to 5 kg
Friction	Uniform 0-1.5
Initial state	Position ±1 m, yaw ±π, base linear and angular velocity uniform ±0.5
Observation noise	Arm joint positions ±0.01 rad, velocities ±1.5 rad/s (student only)
End-effector disturbance	Periodic force and torques described above

Code skeleton

The environment is one class on top of LeggedRobot. Each TODO refers to a section of this page.

# legged_gym/envs/paint/paint_env.py
class PaintEnv(LeggedRobot):
    """Transport task, shared by the teacher task and the student task."""

    def _init_buffers(self):
        super()._init_buffers()
        # TODO: load the exported low-level policy (TorchScript) and its leg-action buffer
        # TODO: buffers for the wrench target, the current wrench c(t), external forces
        #       and torques, and the student observation history
        # TODO: acquire the rigid-body state tensor (payload and end-effector states)

    def _process_rigid_body_props(self, props, env_id):
        # TODO: base mass randomization (super) and payload mass ~ U(0, 10) kg
        ...

    def reset_idx(self, env_ids):
        super().reset_idx(env_ids)
        # TODO: resample the wrench target (zero with probability 0.02)
        # TODO: clear the observation history and the low-level policy buffers

    def step(self, actions):
        # TODO: base command a[0:3] -> limits and deadband -> low-level policy
        #       -> leg joint targets
        # TODO: arm joint targets = default + 0.5 * a[3:9]
        # TODO: for each of the 4 physics steps: PD torques, apply the interaction wrench,
        #       simulate
        # TODO: post_physics_step(), then return obs, privileged obs, reward, reset, extras
        ...

    def _apply_interaction_wrench(self):
        # TODO: c(t) = s(t) * target; force 40 * c_xy [N] and yaw torque 10 * c_z [N m]
        #       in the base frame, rotated to the world frame, applied to the payload body
        # TODO: add the periodic end-effector disturbance
        ...

    def compute_observations(self):
        # TODO: push the student terms into the H-step history (Observations table)
        # TODO: teacher = student history + current wrench;
        #       critic = teacher observation + end-effector and payload states
        # TODO: obs_buf = teacher or student observation (depending on the task),
        #       privileged_obs_buf = critic, extras["teacher_obs"] = teacher observation
        ...

    def _update_terrain_curriculum(self, env_ids):
        # TODO: as in LeggedRobot, with the wrench reference instead of the velocity command
        ...

    # Rewards (Table I). legged_gym provides torques and termination.
    def _reward_force_tracking(self):  # TODO: r_force
        ...

    def _reward_torque_tracking(self):  # TODO: r_torque
        ...

    def _reward_payload_height(self):  # TODO: r_h
        ...

    def _reward_dof_penalty(self):  # TODO: joint DoF penalty
        ...

    def _reward_action_penalty(self):  # TODO: action penalty
        ...

    def _reward_arm_posture(self):  # TODO: arm posture penalty
        ...