Replies: 1 comment 1 reply
|
Hi @btaba, thanks for the reply! It works, I can move the goal successfully now by hard-coding the mocap_pos. My question now is that I am setting the mocap_pos of the goal to a random position in the reset function, it seems to be getting set to exact same value even across runs. I figures this is what domain_randomization is supposed to be used for, but as far as I understand I only have access to sys, but mocap_pos lives in data. How am I supposed to randomize the goal position so the agent learns a more generalizable policy? from datetime import datetime
from etils import epath
import functools
from IPython.display import HTML
from typing import Any, Dict, Sequence, Tuple, Union, Optional
import os
from ml_collections import config_dict
import jax
from jax import numpy as jp
import numpy as np
from flax.training import orbax_utils
from flax import struct
from matplotlib import pyplot as plt
import mediapy as media
from orbax import checkpoint as ocp
import mujoco
from mujoco import mjx
from brax import base
from brax import envs
from brax import math
from brax.base import Base, Motion, Transform
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.training.agents.ppo import train as ppo
from brax.training.agents.ppo import networks as ppo_networks
from brax.io import html, mjcf, model
# Path to car_hazard_goal.xml - adjust if needed
CAR_HAZARD_GOAL_XML_PATH = 'xml/car_hazard_goal.xml'
def default_config() -> config_dict.ConfigDict:
"""Returns the default config for CarHazardGoal environment."""
config = config_dict.create(
# New safety-gymnasium reward parameters
reward_distance=1.0, # Dense reward scale for distance moved to the goal
reward_goal=1.0, # Sparse reward for reaching the goal
goal_size=0.7, # Distance threshold for achieving the goal
reward_orientation=False, # Optional: Reward for maintaining upright orientation
reward_orientation_scale=0.002, # Scale for orientation reward
reward_orientation_body='agent', # Body to check orientation (unused if reward_orientation=False)
# Other parameters (kept or adjusted)
terminate_when_unhealthy=True, # Keep termination based on health
healthy_z_range=(0.05, 0.3), # Keep health definition
reset_noise_scale=0.005,
exclude_current_positions_from_observation=True,
max_velocity=5.0, # Keep velocity limit for calculation stability
debug=False,
)
return config
def safe_norm(x, axis=None, keepdims=False, eps=1e-8):
"""Safely compute the norm with a small epsilon to avoid NaN."""
return jp.sqrt(jp.sum(jp.square(x), axis=axis, keepdims=keepdims) + eps)
# Add JAX-compatible helper function for NaN handling
def nan_to_zero(x):
"""Replace NaN values with zeros in a JAX-compatible way."""
return jp.where(jp.isnan(x), jp.zeros_like(x), x)
class CarHazardGoal(PipelineEnv):
def __init__(
self,
config: config_dict.ConfigDict = default_config(),
config_overrides: Optional[Dict[str, Any]] = None,
**kwargs,
):
# Load the car model from XML
mj_model = mujoco.MjModel.from_xml_path(CAR_HAZARD_GOAL_XML_PATH)
mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG
mj_model.opt.iterations = 4
mj_model.opt.ls_iterations = 4
# Get body IDs directly from MuJoCo model before loading into Brax
try:
self._agent_body_name = 'agent' # Store name for orientation reward if needed
self._agent_body = mj_model.body(self._agent_body_name).id
self._goal_body = mj_model.body('goal').id # Still useful for observation? Or use mocap? Mocap is used later.
# Get hazard body IDs (kept for potential future use, but not reward)
self._hazard_bodies = []
for i in range(1, 5):
hazard_name = f'hazard{i}'
try:
self._hazard_bodies.append(mj_model.body(hazard_name).id)
except Exception:
pass # Skip if hazard doesn't exist
except Exception as e:
print(f"Warning: Error getting body IDs: {e}")
# Fallback to index-based approach
self._agent_body_name = 'agent'
self._agent_body = 1
self._goal_body = 2
self._hazard_bodies = [3, 4, 5, 6]
sys = mjcf.load_model(mj_model)
# Find the mocap ID for the goal body
self._goal_mocap_id = None
if mj_model.nmocap > 0:
for i in range(mj_model.nbody):
if mj_model.body(i).name == "goal" and mj_model.body(i).mocapid >= 0:
print(f"Goal mocap ID found: {mj_model.body(i).mocapid}")
self._goal_mocap_id = mj_model.body(i).mocapid
break
else:
# Fallback if 'goal' mocap body not found by name
if mj_model.nmocap > 0:
print("Warning: 'goal' body with mocapid >= 0 not found. Defaulting to mocapid 0.")
self._goal_mocap_id = 0
else:
print("Error: No mocap bodies found in the model.")
# Or raise an error, depending on requirements
self._goal_mocap_id = -1 # Indicate error or invalid state
else:
print("Error: No mocap bodies defined in the model.")
self._goal_mocap_id = -1 # Indicate error or invalid state
physics_steps_per_control_step = 4
kwargs['n_frames'] = kwargs.get(
'n_frames', physics_steps_per_control_step)
kwargs['backend'] = 'mjx'
super().__init__(sys, **kwargs)
# Apply config overrides if provided
if config_overrides:
config = config.copy_and_resolve_references()
for k, v in config_overrides.items():
config[k] = v
self._config = config
# Store new reward parameters
self._reward_distance = config.reward_distance
self._reward_goal = config.reward_goal
self._goal_size = config.goal_size
self._reward_orientation = config.reward_orientation
self._reward_orientation_scale = config.reward_orientation_scale
self._reward_orientation_body = config.reward_orientation_body # Name stored
# Store other necessary parameters
self._terminate_when_unhealthy = config.terminate_when_unhealthy
self._healthy_z_range = config.healthy_z_range
self._reset_noise_scale = config.reset_noise_scale
self._exclude_current_positions_from_observation = (
config.exclude_current_positions_from_observation
)
self._max_velocity = config.max_velocity
self._debug = config.debug
def reset(self, rng: jp.ndarray) -> State:
"""Resets the environment to an initial state with randomized goal."""
rng, rng1, rng2, rng_goal = jax.random.split(rng, 4)
# Randomize initial car position with small noise
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
)
# Ensure qpos has valid quaternion
if qpos.shape[0] > 6: # If quaternion exists
quat_norm = safe_norm(qpos[3:7])
qpos = qpos.at[3:7].set(qpos[3:7] / quat_norm)
# Initialize environment
data = self.pipeline_init(qpos, qvel)
# Randomize goal position (with more moderate bounds)
goal_pos = jax.random.uniform(
rng_goal,
(3,),
minval=jp.array([-2.0, -2.0, 0.09]),
maxval=jp.array([2.0, 2.0, 0.09]),
)
# Get positions
agent_pos = data.xpos[self._agent_body]
# Set goal position in the identified mocap body
# Add check to ensure _goal_mocap_id is valid
if self._goal_mocap_id is not None and self._goal_mocap_id >= 0:
data = data.replace(mocap_pos=data.mocap_pos.at[self._goal_mocap_id].set(goal_pos))
print(f"Goal position set to: {goal_pos}")
else:
print("Error: Invalid or missing goal mocap ID. Cannot set goal position.")
# Handle error case: maybe set done=True or raise exception
# For now, we proceed but the goal won't be set correctly
pass
# Calculate initial distance to goal for the first step
initial_dist_goal = safe_norm(agent_pos[:2] - goal_pos[:2])
# Store goal position and initial distance for reward calculation
info = {
"goal_pos": goal_pos,
"step_count": 0,
"last_obs": None,
"last_dist_goal": initial_dist_goal, # Initialize last_dist_goal
}
# Get observation
obs = self._get_obs(data)
# Store observation in info for stability check
info["last_obs"] = obs
reward, done, zero = jp.zeros(3)
metrics = {
# New reward components
'dist_reward': zero, # Reward from distance change
'goal_reward': zero, # Reward for reaching goal
'orientation_reward': zero, # Reward for orientation (if enabled)
'reward': zero, # Total reward (expected by wrappers)
# Other potentially useful metrics
'x_position': jp.array(data.xpos[self._agent_body, 0]),
'y_position': jp.array(data.xpos[self._agent_body, 1]),
'distance_to_goal': initial_dist_goal,
'x_velocity': zero,
'y_velocity': zero,
'z_alignment': zero, # Z-alignment metric (even if reward disabled)
}
return State(data, obs, reward, done, metrics, info)
def step(self, state: State, action: jp.ndarray) -> State:
"""Runs one timestep of the environment's dynamics."""
# Increment step counter
step_count = state.info.get('step_count', 0) + 1
# Get the last valid observation and distance
last_obs = state.info.get('last_obs', state.obs)
last_dist_goal = state.info['last_dist_goal']
# Clip action for stability (optional, can be removed if not needed)
action = jp.clip(action, -1.0, 1.0)
data0 = state.pipeline_state
data = self.pipeline_step(data0, action)
# Get positions
agent_pos = data.xpos[self._agent_body]
goal_pos = state.info['goal_pos']
# --- Calculate Reward Components ---
# 1. Distance-Based Reward
dist_goal = safe_norm(agent_pos[:2] - goal_pos[:2])
dist_reward = (last_dist_goal - dist_goal) * self._reward_distance
# 2. Goal Achievement Reward
goal_achieved = dist_goal <= self._goal_size
goal_reward = jp.where(goal_achieved, self._reward_goal, 0.0)
# 3. Orientation Reward (Optional)
z_alignment = 0.0
orientation_reward = 0.0
if self._reward_orientation:
# Get quaternion of the agent body (assuming it's the first part of qpos)
# Check if the orientation body exists in the model
# Note: This assumes the agent's quaternion is at the start of qpos after potential exclusion.
# Adjust indexing if necessary based on your MJCF structure and observation exclusion.
quat = data.qpos[3:7] # Assuming standard floating base: [x, y, z, qw, qx, qy, qz, ...] -> qw, qx, qy, qz indices 3-7
# Normalize quaternion
quat_norm = safe_norm(quat)
quat = quat / jp.maximum(quat_norm, 1e-8)
# Calculate rotation matrix
rot_matrix = math.quat_to_rot(quat)
# zalign = R[2, 2]
z_alignment = rot_matrix[2, 2]
orientation_reward = self._reward_orientation_scale * z_alignment
# --- Total Reward ---
reward = dist_reward + goal_reward + orientation_reward
# --- Health Check (for termination, not reward) ---
min_z, max_z = self._healthy_z_range
is_healthy = jp.logical_and(
agent_pos[2] >= min_z,
agent_pos[2] <= max_z
).astype(jp.float32)
# --- Termination Conditions ---
# Check for NaN in state using JAX-compatible operations
# Calculate velocity for NaN check (keep calculation local if only needed here)
dt = jp.maximum(self.dt, 1e-6) # Ensure dt is not too small
velocity = (agent_pos - data0.xpos[self._agent_body]) / dt
has_nan_pos = jp.any(jp.isnan(agent_pos))
has_nan_vel = jp.any(jp.isnan(velocity))
has_nan_state = jp.logical_or(has_nan_pos, has_nan_vel)
# Terminate if unhealthy (optional), goal reached, or NaN state
done = jp.logical_or(
jp.logical_or(
(1.0 - is_healthy) * self._terminate_when_unhealthy,
goal_achieved
),
has_nan_state
)
# --- Observation and State Update ---
obs = self._get_obs(data)
# Handle NaN observation using JAX-compatible operations
has_nan_obs = jp.any(jp.isnan(obs))
obs = jp.where(has_nan_obs, last_obs, obs)
# Update info dictionary - update existing dict instead of replacing
# Ensure all keys from input state.info are preserved
new_info = state.info.copy() # Start with a copy of the input info
new_info.update({
"goal_pos": goal_pos, # Update goal position (usually static but good practice)
"step_count": step_count, # Update internal step counter if needed
"last_obs": obs, # Store current observation for next step's fallback
"last_dist_goal": dist_goal, # Store current distance for next step's reward calculation
})
# Update metrics safely
metrics = {
'dist_reward': dist_reward,
'goal_reward': goal_reward,
'orientation_reward': orientation_reward,
'reward': reward, # Add total reward to metrics
'x_position': jp.clip(agent_pos[0], -10.0, 10.0), # Keep clipping for metric stability
'y_position': jp.clip(agent_pos[1], -10.0, 10.0),
'distance_to_goal': dist_goal,
'x_velocity': jp.clip(velocity[0], -self._max_velocity, self._max_velocity),
'y_velocity': jp.clip(velocity[1], -self._max_velocity, self._max_velocity),
'z_alignment': z_alignment, # Log z_alignment
}
# Create fresh State
return State(data, obs, reward, done.astype(jp.float32), metrics, new_info)
def _get_obs(self, data: mjx.Data) -> jp.ndarray:
"""Creates an observation that includes car state and goal information."""
agent_pos = data.xpos[self._agent_body]
# Get goal position from the identified mocap body
# Add check for valid mocap ID
if self._goal_mocap_id is not None and self._goal_mocap_id >= 0:
goal_pos = data.mocap_pos[self._goal_mocap_id]
else:
# Handle error: Use a default or raise an error
print("Warning: Invalid or missing goal mocap ID in _get_obs. Using origin as default goal.")
goal_pos = jp.zeros(3)
position = data.qpos
if self._exclude_current_positions_from_observation:
position = position[2:] # Exclude global x,y position
# Include relevant state information with clipping for stability
qpos_obs = jp.clip(position, -10.0, 10.0)
qvel_obs = jp.clip(data.qvel, -10.0, 10.0)
# Relative position to goal with checks
rel_goal_pos = jp.clip(goal_pos - agent_pos, -10.0, 10.0)
# Squeeze potential extra dimension from rel_goal_pos
rel_goal_pos = jp.squeeze(rel_goal_pos)
# Orientation - ensure it's properly normalized
if position.shape[0] >= 4: # If quaternion exists
orientation = position[:4]
quat_norm = safe_norm(orientation)
orientation = orientation / jp.maximum(quat_norm, 1e-6)
else:
orientation = jp.array([1.0, 0.0, 0.0, 0.0]) # Default quaternion
# Actuator forces
actuator_forces = jp.clip(data.qfrc_actuator, -10.0, 10.0)
# Build observation with JAX-compatible NaN handling
obs = jp.concatenate([
nan_to_zero(qpos_obs),
nan_to_zero(qvel_obs),
nan_to_zero(rel_goal_pos),
nan_to_zero(orientation),
nan_to_zero(actuator_forces),
])
return obs
# Register the environment
envs.register_environment('car_hazard_goal', CarHazardGoal) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Uh oh!
There was an error while loading. Please reload this page.
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
All reactions