Hi, I'm working on Brax for my master thesis in Safe Reinforcement Learning, where I want to integrate common tasks into Brax. One of these tasks is for the agent to walk to a goal while avoiding hazards. I would like the goal to change position after the agent reaches it. However, I'm struggling to get it actually rendering the new position of the goal.
humanoid_hazard_goal.xml
<mujoco model="humanoid_hazard">
<include file="humanoid.xml"/>
<worldbody>
<!-- Goal (green box) -->
<body name="goal" pos="2 2 0.09">
<geom type="box" name="goal" size="0.5 0.5 0.6" condim="3"
friction="1 .03 .003" rgba="0 1 0 1" contype="2" conaffinity="1" solref="0.01 1"/>
</body>
<!-- Hazards (red boxes) -->
<body name="hazard1" pos="-2.4 2.4 0.09">
<geom type="box" name="hazard1" size="0.5 0.5 0.6" condim="3"
friction="1 .03 .003" rgba="1 0 0 1" contype="2" conaffinity="1" solref="0.01 1"/>
</body>
<body name="hazard2" pos="2.4 -2.4 0.09">
<geom type="box" name="hazard2" size="0.5 0.5 0.6" condim="3"
friction="1 .03 .003" rgba="1 0 0 1" contype="2" conaffinity="1" solref="0.01 1"/>
</body>
<body name="hazard3" pos="-1.8 -2.4 0.09">
<geom type="box" name="hazard3" size="0.5 0.5 0.6" condim="3"
friction="1 .03 .003" rgba="1 0 0 1" contype="2" conaffinity="1" solref="0.01 1"/>
</body>
</worldbody>
</mujoco>I've created various environments, but for testing purposes, I'm trying to make it work with this environment where I set a new position for the goal every 10 steps.
import jax
import jax.numpy as jp
import mujoco
import os
from brax import base
from brax import envs
from brax.base import State as PipelineState
from brax.envs.base import Env, PipelineEnv, State
from brax.mjx.base import State as MjxState
from brax.io import mjcf
from typing import Tuple
class SimpleResettingGoal(PipelineEnv):
"""Environment with goal respawning every 10 steps."""
def __init__(
self,
forward_reward_weight=1.25,
ctrl_cost_weight=0.1,
healthy_reward=5.0,
terminate_when_unhealthy=True,
healthy_z_range=(1.0, 2.0),
reset_noise_scale=1e-2,
exclude_current_positions_from_observation=True,
goal_reset_steps=10, # Parameter for goal reset frequency
goal_area=(-5.0, -5.0, 5.0, 5.0), # xmin, ymin, xmax, ymax
**kwargs,
):
# Get XML path
current_dir = os.path.dirname(os.path.abspath(__file__))
xml_path = os.path.join(os.path.dirname(current_dir), 'xml', 'humanoid_hazard_goal.xml')
mj_model = mujoco.MjModel.from_xml_path(xml_path)
mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG
mj_model.opt.iterations = 6
mj_model.opt.ls_iterations = 6
sys = mjcf.load_model(mj_model)
physics_steps_per_control_step = 5
kwargs['n_frames'] = kwargs.get(
'n_frames', physics_steps_per_control_step)
kwargs['backend'] = 'mjx'
super().__init__(sys, **kwargs)
self._forward_reward_weight = forward_reward_weight
self._ctrl_cost_weight = ctrl_cost_weight
self._healthy_reward = healthy_reward
self._terminate_when_unhealthy = terminate_when_unhealthy
self._healthy_z_range = healthy_z_range
self._reset_noise_scale = reset_noise_scale
self._exclude_current_positions_from_observation = (
exclude_current_positions_from_observation
)
self._goal_body = mj_model.body('goal').id
self._hazard_body_1 = mj_model.body('hazard1').id
self._hazard_body_2 = mj_model.body('hazard2').id
self._hazard_body_3 = mj_model.body('hazard3').id
self.goal_reset_steps = goal_reset_steps # Store the reset frequency
self.goal_area = goal_area # Store goal area
def _generate_random_goal_position(self, rng):
"""Generate a random goal position within the defined area."""
xmin, ymin, xmax, ymax = self.goal_area
goal_x = jax.random.uniform(rng, (), minval=xmin, maxval=xmax)
goal_y = jax.random.uniform(rng, (), minval=ymin, maxval=ymax)
goal_pos = jp.array([goal_x, goal_y, 1.0]) # Fixed z-coordinate
return goal_pos
def reset(self, rng: jp.ndarray) -> State:
"""Resets the environment to an initial state."""
rng, rng1, rng2, rng3 = jax.random.split(rng, 4)
low, hi = -self._reset_noise_scale, self._reset_noise_scale
qpos = self.sys.qpos0 + jax.random.uniform(
rng1, (self.sys.nq,), minval=low, maxval=hi
)
qvel = jax.random.uniform(
rng2, (self.sys.nv,), minval=low, maxval=hi
)
# Initial placement of the goal (using _generate_random_goal_position)
goal_pos = self._generate_random_goal_position(rng3)
pipeline_state = self.pipeline_init(qpos, qvel)
pipeline_state = pipeline_state.replace(
xpos=pipeline_state.xpos.at[self._goal_body].set(goal_pos) # Set initial goal position.
)
obs = self._get_obs(pipeline_state)
reward, done, zero = jp.zeros(3)
metrics = {
'reward_forward': zero,
'reward_ctrl': zero,
'reward_healthy': zero,
'x_position': zero,
'y_position': zero,
}
info = {
'goal_pos': goal_pos, # Store for later use.
'steps': 0,
'goal_resets': 0
} # Initialize step counter
return State(pipeline_state, obs, reward, done, metrics, info)
def step(self, state: State, action: jp.ndarray) -> State:
"""Run one timestep of the environment's dynamics."""
data = self.pipeline_step(state.pipeline_state, action)
# --- Goal Reset Logic ---
goal_pos = state.info['goal_pos']
steps = state.info['steps'] + 1
goal_resets = state.info['goal_resets']
rng = jax.random.PRNGKey(steps)
should_reset_goal = (steps % self.goal_reset_steps) == 0
def true_fn(_):
new_rng, rng1 = jax.random.split(rng)
new_goal_pos = self._generate_random_goal_position(rng1)
new_data = data.replace(xpos=data.xpos.at[self._goal_body].set(new_goal_pos))
return new_rng, new_data, new_goal_pos, goal_resets + 1
def false_fn(_):
return rng, data, goal_pos, goal_resets
# Apply the conditional
rng, data, goal_pos, goal_resets = jax.lax.cond(
should_reset_goal,
true_fn,
false_fn,
None # Dummy operand that gets passed to both branches
)
# --- Reward Calculation (Example) ---
# Replace with your actual reward/cost function
reward = 1.0 # Placeholder: Give a small reward for existing.
done = jp.array(0, dtype=jp.float32)
# --- Observation ---
obs = self._get_obs(data)
# --- Update State ---
new_info = state.info.copy()
new_info.update({
'goal_pos': goal_pos,
'steps': steps,
'goal_resets': goal_resets
}) # Update step count in info
return state.replace(
pipeline_state=data, obs=obs, reward=reward, done=done, info=new_info
)
def _get_obs(self, data: MjxState) -> jp.ndarray:
"""Return relevant observations."""
# Example: Include goal position in observation. Adapt as needed.
qpos = data.qpos
qvel = data.qvel
# Get the goal position from the body
goal_pos = data.xpos[self._goal_body, :2]
# Get the humanoid position from the com
humanoid_pos = data.xpos[1, :2]
# Relative position
relative_pos = goal_pos - humanoid_pos
# Concatenate position and velocity of the humanoid, and add the relative position
obs = jp.concatenate([qpos[2:], qvel, relative_pos]) # Exclude root z-height.
return obsI then use this code to render the rollout:
# instantiate the environment
env_name = 'SimpleResettingGoal'
env = envs.get_environment(env_name)
# define the jit reset/step functions
jit_reset = jax.jit(env.reset)
jit_step = jax.jit(env.step)
# initialize the state
state = jit_reset(jax.random.PRNGKey(0))
rollout = [state.pipeline_state] # store the pipeline_state for rendering
# grab a trajectory
first = True
for i in range(200):
ctrl = -0.1 * jp.ones(env.sys.nu)
state = jit_step(state, ctrl)
rollout.append(state.pipeline_state) # append pipeline_state, not whole state
first = False
media.show_video(env.render(rollout, camera='side'), fps=1.0 / env.dt)
In the video, the goal just stays at the initial position given in the xml. Please let me know how I can manage to do this and whether I need to provide any more information
Hi, I'm working on Brax for my master thesis in Safe Reinforcement Learning, where I want to integrate common tasks into Brax. One of these tasks is for the agent to walk to a goal while avoiding hazards. I would like the goal to change position after the agent reaches it. However, I'm struggling to get it actually rendering the new position of the goal.
humanoid_hazard_goal.xml
I've created various environments, but for testing purposes, I'm trying to make it work with this environment where I set a new position for the goal every 10 steps.
I then use this code to render the rollout:
In the video, the goal just stays at the initial position given in the xml. Please let me know how I can manage to do this and whether I need to provide any more information