From a2d413002b1092e2d2fcc0027e042e2c6250c601 Mon Sep 17 00:00:00 2001 From: Ruofan Kong Date: Tue, 5 May 2020 16:27:40 -0700 Subject: [PATCH] No Case: SAC-DisCor. --- python/ray/rllib/agents/sac/sac_model.py | 30 +- .../ray/rllib/agents/sac/sac_policy_graph.py | 71 +- python/ray/rllib/tests/moab_env/__init__.py | 15 + python/ray/rllib/tests/moab_env/moab_env.py | 309 +++++++ python/ray/rllib/tests/moab_env/moab_model.py | 773 ++++++++++++++++++ python/ray/rllib/tests/test_convergence.py | 52 +- 6 files changed, 1236 insertions(+), 14 deletions(-) create mode 100644 python/ray/rllib/tests/moab_env/__init__.py create mode 100644 python/ray/rllib/tests/moab_env/moab_env.py create mode 100644 python/ray/rllib/tests/moab_env/moab_model.py diff --git a/python/ray/rllib/agents/sac/sac_model.py b/python/ray/rllib/agents/sac/sac_model.py index 57128304c987..d558f368bbfc 100644 --- a/python/ray/rllib/agents/sac/sac_model.py +++ b/python/ray/rllib/agents/sac/sac_model.py @@ -193,9 +193,25 @@ def build_q_net(name, observations, actions): q_net(dict(observations=observations, actions=actions)) ) + def build_error_net(name, observations, actions): + q_net = build_fcn( + input_shapes=dict(observations=(num_outputs,), actions=(self.action_dim,)), + num_outputs=1, + hidden_layer_sizes=(256, 256, 256), + hidden_activations="relu", + name=name + ) + return tf.keras.Model( + [observations, actions], + q_net(dict(observations=observations, actions=actions)) + ) + self.q_net = build_q_net("q", self.model_out, self.actions) self.register_variables(self.q_net.variables) + self.error_net = build_error_net("error", self.model_out, self.actions) + self.register_variables(self.error_net.variables) + if twin_q: self.twin_q_net = build_q_net("twin_q", self.model_out, self.actions) self.register_variables(self.twin_q_net.variables) @@ -203,9 +219,10 @@ def build_q_net(name, observations, actions): self.twin_q_net = None self.log_alpha = tf.Variable(0.0, dtype=tf.float32, name="log_alpha") + self.tau_temp = tf.Variable(10.0, dtype=tf.float32, name="tau_temp") self.alpha = tf.exp(self.log_alpha) - self.register_variables([self.log_alpha]) + self.register_variables([self.tau_temp, self.log_alpha]) def get_policy_output(self, model_out, deterministic=False): """Return the (unscaled) output of the policy network. @@ -273,3 +290,14 @@ def q_variables(self): return self.q_net.variables + ( self.twin_q_net.variables if self.twin_q_net else [] ) + + def error_values(self, model_out, actions): + """ + Return the error estimates for the most recent forward pass. + This implements Error(s, a). + """ + return self.error_net([model_out, actions]) + + def error_variables(self): + """Return the list of variables for Error nets.""" + return list(self.error_net.variables) diff --git a/python/ray/rllib/agents/sac/sac_policy_graph.py b/python/ray/rllib/agents/sac/sac_policy_graph.py index 9724aba6cbe2..de58192aad7e 100644 --- a/python/ray/rllib/agents/sac/sac_policy_graph.py +++ b/python/ray/rllib/agents/sac/sac_policy_graph.py @@ -167,6 +167,7 @@ def actor_critic_loss(policy, model, _, train_batch): log_alpha = model.log_alpha alpha = model.alpha + tau_temp = model.tau_temp # q network evaluation q_t = model.get_q_values(model_out_t, train_batch[SampleBatch.ACTIONS]) @@ -207,6 +208,18 @@ def actor_critic_loss(policy, model, _, train_batch): + policy.config["gamma"] ** policy.config["n_step"] * q_tp1_best_masked ) + # Compute errors and target errors at next state, and an action from the policy. + qf_pred_errs = model.error_values(model_out_tp1, policy_tp1) + qf_pred_errs_t = policy.target_model.error_values(model_out_tp1, policy_tp1) + + # Moving mean of the error values over batches. + err_logits = - tf.stop_gradient( + policy.config["gamma"] * qf_pred_errs / tau_temp + ) + squeezed_err_logits = tf.squeeze(err_logits, axis=len(err_logits.shape) - 1) + + err_values = model.error_values(model_out_t, policy_t) + # compute the error (potentially clipped) if policy.config["twin_q"]: base_td_error = q_t_selected - q_t_selected_target @@ -215,15 +228,32 @@ def actor_critic_loss(policy, model, _, train_batch): else: td_error = tf.square(q_t_selected - q_t_selected_target) + err_targets = tf.stop_gradient( + tf.abs(td_error) + (policy.config["gamma"] * qf_pred_errs_t)) + + # This is used to update the moving mean, self._error_model_tau_ph + mean_error_values = tf.reduce_mean(err_values) + tau_temp_update_ops = tau_temp.assign( + mean_error_values * policy.config["tau"] + ( + 1.0 - policy.config["tau"]) * tau_temp + ) + error_loss = tf.losses.mean_squared_error( + labels=err_targets, predictions=err_values, weights=0.5 + ) + critic_loss = [ tf.losses.mean_squared_error( - labels=q_t_selected_target, predictions=q_t_selected, weights=0.5 + labels=squeezed_err_logits * q_t_selected_target, + predictions=squeezed_err_logits * q_t_selected, + weights=0.5 ) ] if policy.config["twin_q"]: critic_loss.append( tf.losses.mean_squared_error( - labels=q_t_selected_target, predictions=twin_q_t_selected, weights=0.5 + labels=squeezed_err_logits * q_t_selected_target, + predictions=squeezed_err_logits * twin_q_t_selected, + weights=0.5 ) ) @@ -243,14 +273,22 @@ def actor_critic_loss(policy, model, _, train_batch): policy.actor_loss = actor_loss policy.critic_loss = critic_loss policy.alpha_loss = alpha_loss + policy.error_loss = error_loss + policy.tau_temp_update_ops = tau_temp_update_ops # in a custom apply op we handle the losses separately, but return them # combined in one loss for now - return actor_loss + tf.add_n(critic_loss) + alpha_loss + return error_loss + actor_loss + tf.add_n(critic_loss) + alpha_loss def gradients(policy, optimizer, loss): if policy.config["grad_norm_clipping"] is not None: + error_grads_and_vars = _minimize_and_clip( + optimizer, + policy.error_loss, + var_list=policy.model.error_variables(), + clip_val=policy.config["grad_norm_clipping"], + ) actor_grads_and_vars = _minimize_and_clip( optimizer, policy.actor_loss, @@ -287,6 +325,10 @@ def gradients(policy, optimizer, loss): clip_val=policy.config["grad_norm_clipping"], ) else: + error_grads_and_vars = policy._error_optimizer.compute_gradients( + policy.error_loss, + var_list=policy.model.error_variables() + ) actor_grads_and_vars = policy._actor_optimizer.compute_gradients( policy.actor_loss, var_list=policy.model.policy_variables() ) @@ -308,6 +350,9 @@ def gradients(policy, optimizer, loss): ) # save these for later use in build_apply_op + policy._error_grads_and_vars = [ + (g, v) for (g, v) in error_grads_and_vars if g is not None + ] policy._actor_grads_and_vars = [ (g, v) for (g, v) in actor_grads_and_vars if g is not None ] @@ -318,7 +363,8 @@ def gradients(policy, optimizer, loss): (g, v) for (g, v) in alpha_grads_and_vars if g is not None ] grads_and_vars = ( - policy._actor_grads_and_vars + policy._error_grads_and_vars + + policy._actor_grads_and_vars + policy._critic_grads_and_vars + policy._alpha_grads_and_vars ) @@ -326,6 +372,10 @@ def gradients(policy, optimizer, loss): def apply_gradients(policy, optimizer, grads_and_vars): + discor_steps = tf.Variable(1, name='discor_steps', trainable=False) + error_apply_ops = policy._error_optimizer.apply_gradients( + policy._error_grads_and_vars, global_step=discor_steps + ) actor_apply_ops = policy._actor_optimizer.apply_gradients( policy._actor_grads_and_vars ) @@ -334,16 +384,18 @@ def apply_gradients(policy, optimizer, grads_and_vars): half_cutoff = len(cgrads) // 2 if policy.config["twin_q"]: critic_apply_ops = [ - policy._critic_optimizer[0].apply_gradients(cgrads[:half_cutoff]), - policy._critic_optimizer[1].apply_gradients(cgrads[half_cutoff:]), + policy._critic_optimizer[0].apply_gradients(cgrads[:half_cutoff], global_step=discor_steps), + policy._critic_optimizer[1].apply_gradients(cgrads[half_cutoff:], global_step=discor_steps), ] else: - critic_apply_ops = [policy._critic_optimizer[0].apply_gradients(cgrads)] + critic_apply_ops = [policy._critic_optimizer[0].apply_gradients(cgrads, global_step=discor_steps)] alpha_apply_ops = policy._alpha_optimizer.apply_gradients( policy._alpha_grads_and_vars, global_step=tf.train.get_or_create_global_step() ) - return tf.group([actor_apply_ops, alpha_apply_ops] + critic_apply_ops) + return tf.group( + [error_apply_ops, actor_apply_ops, alpha_apply_ops + ] + critic_apply_ops) def stats(policy, train_batch): @@ -351,9 +403,11 @@ def stats(policy, train_batch): "td_error": tf.reduce_mean(policy.td_error), "actor_loss": tf.reduce_mean(policy.actor_loss), "critic_loss": tf.reduce_mean(policy.critic_loss), + "error_loss": tf.reduce_mean(policy.error_loss), "mean_q": tf.reduce_mean(policy.q_t), "max_q": tf.reduce_max(policy.q_t), "min_q": tf.reduce_min(policy.q_t), + "update_tau_temp_ops": policy.tau_temp_update_ops } @@ -377,6 +431,7 @@ def __init__(self, config): self.global_step = tf.train.get_or_create_global_step() # use separate optimizers for actor & critic + self._error_optimizer = tf.train.AdamOptimizer(learning_rate=3e-4) self._actor_optimizer = tf.train.AdamOptimizer( learning_rate=config["optimization"]["actor_learning_rate"] ) diff --git a/python/ray/rllib/tests/moab_env/__init__.py b/python/ray/rllib/tests/moab_env/__init__.py new file mode 100644 index 000000000000..4725977b91d2 --- /dev/null +++ b/python/ray/rllib/tests/moab_env/__init__.py @@ -0,0 +1,15 @@ +from gym.envs.registration import register + +# Register a custom environment for frozen lake with deterministic states +register( + id='Moab-v0', + entry_point='moab_env.moab_env:MoabSim', + kwargs={ + "config":{ + "use_normalize_action": True, + "use_normalize_state": True, + "use_dr": False + } + }, + reward_threshold=1000, +) diff --git a/python/ray/rllib/tests/moab_env/moab_env.py b/python/ray/rllib/tests/moab_env/moab_env.py new file mode 100644 index 000000000000..372100186548 --- /dev/null +++ b/python/ray/rllib/tests/moab_env/moab_env.py @@ -0,0 +1,309 @@ +import sys +from typing import Optional, Tuple + +import gym +import numpy as np +from gym.spaces import Box +from pyrr import vector, matrix33 + +from moab_env.moab_model import MoabModel, clamp + + +_CLOSE_ENOUGH = 0.01 +_DEFAULT_BALL_NOISE = 0.0005 +_DEFAULT_FRICTION = 0.6 +_DEFAULT_GRAVITY = 9.81 +_DEFAULT_HEIGHT_LIMIT = 0.0 +_DEFAULT_PLATE_NOISE = (np.pi / 180.0) * 1 +_DEFAULT_TILT_LIMIT = (np.pi / 180.0) * 15 +_DEFAULT_TIME_DELTA = 0.02 +_MAX_ITER_COUNT = 250 +_MAX_FRICTION = 1.0 +_MAX_GRAVITY = _DEFAULT_GRAVITY * 2 +_MAX_TILT_VELOCITY = (np.pi / 180.0) * 300 +_MAX_VELOCITY = 1.0 +_MIN_FRICTION = 0.2 +_MIN_GRAVITY = _DEFAULT_GRAVITY / 2 +_PING_PONG_COR = 0.89 +_PING_PONG_MASS = 0.0027 +_PING_PONG_RADIUS = 0.020 +_PING_PONG_SHELL = 0.0002 +_PLATE_RADIUS = 0.225 / 2.0 +_PLATE_Z_OFFSET = 0.009 +_TILT_ACC = (60.0 / 3.0) * _MAX_TILT_VELOCITY + + +class MoabSim(gym.Env): + + """ + Arguments: + config: needs to incorporate the following parameters + seed: The random seed for the random number generator. + use_dr: The flag to turn on the domain randomization. + use_normalize_action: The flag to enable the action normalization. + Normally, we use it for PPO. For SAC and other + algorithms, RLLib may take care of it. + use_normalize_state: The flag to enable the range state normalization. + + + Gym-like Moab Simulator - Example: + + >> env = MoabSim(config=dict(seed=None, use_normalize_action=False)) + >> env.reset() + >> episode_reward = 0 + >> episode_count = 0 + >> for _ in range(1000): + >> env.render() + >> _, reward, done, _ = env.step(env.action_space.sample()) + >> episode_reward += reward + >> if done: + >> episode_count += 1 + >> env.reset() + >> print(f"Episode {episode_count} reward: {episode_reward}") + >> episode_reward = 0 + >> env.close() + """ + + def __init__( + self, + config: dict + ): + self.model = MoabModel() + self.model.reset() + self._random_state = np.random.RandomState(config.get("seed")) + self._use_dr = config.get("use_dr", False) + self._use_normalize_action = config.get("use_normalize_action", False) + self._use_normalize_state = config.get("use_normalize_state", True) + print(self._use_dr, self._use_normalize_action, self._use_normalize_state) + + obs_high = np.array( + [sys.maxsize, sys.maxsize, sys.maxsize, sys.maxsize] + ) + obs_low = np.array( + [- sys.maxsize - 1, - sys.maxsize - 1, + - sys.maxsize - 1, - sys.maxsize - 1] + ) + action_high = np.array([1.0, 1.0]) + + self.action_space = Box( + low=-action_high, + high=action_high, + dtype=np.float32 + ) + self.observation_space = Box( + low=obs_low, + high=obs_high, + dtype=np.float32 + ) + + def seed(self, seed: Optional[int] = None) -> None: + """ + Create the random number generator with seed. + """ + self._random_state = np.random.RandomState(seed) + + def step( + self, + action: np.ndarray + ) -> Tuple[np.ndarray, float, bool, dict]: + """ + Perform the one simulation step for Moab simulator. + """ + if self._use_normalize_action: + self.model.roll = clamp( + self._normalize_action(action[0]), -1.0, 1.0) + self.model.pitch = clamp( + self._normalize_action(action[1]), -1.0, 1.0) + else: + self.model.roll = clamp(action[0], -1.0, 1.0) + self.model.pitch = clamp(action[1], -1.0, 1.0) + self.model.step() + return ( + self._get_obs_from_state(), + self._get_reward(), + self._is_penalty_terminal(), + {} + ) + + def reset(self) -> np.ndarray: + """ + Reset the Moab simulator. + """ + return self._do_reset() + + def render( + self, + mode: str = 'human' + ) -> None: + # TODO: Create the visual simulator. + pass + + @staticmethod + def _normalize_action(action: float) -> float: + """ + Normalize the action following the PDP2 output preprocessor. + """ + # NOTE: This follows the action normalizer in PDP2 output + # preprocessor. We suggest using it only when RLLib doesn't have + # the action normalizer for the algorithm. + return 0.5 * np.clip(action, -2.0, 2.0) + + def _do_reset(self) -> np.ndarray: + """ + Reset the simulator with all initial configs depending on the + domain randomization. + """ + self.model.reset() + if self._use_dr: + self.model.roll = self._random_state.uniform(low=-0.2, high=0.2) + self.model.pitch = self._random_state.uniform(low=-0.2, high=0.2) + else: + self.model.roll = -0.1 + self.model.pitch = -0.1 + self.model.height_z = 0.0 + self.model.time_delta = _DEFAULT_TIME_DELTA + self.model.gravity = _DEFAULT_GRAVITY + self.model.friction = _DEFAULT_FRICTION + self.model.plate_radius = _PLATE_RADIUS + self.model.tilt_max_vel = _MAX_TILT_VELOCITY + self.model.tilt_acc = _TILT_ACC + self.model.tilt_limit = _DEFAULT_TILT_LIMIT + self.model.height_z_limit = _DEFAULT_HEIGHT_LIMIT + self.model.ball_mass = _PING_PONG_MASS + self.model.ball_radius = _PING_PONG_RADIUS + self.model.ball_shell = _PING_PONG_SHELL + self.model.ball_COR = _PING_PONG_COR + self.model.target_pos_x = 0.0 + self.model.target_pos_y = 0.0 + self.model.ball_noise = _DEFAULT_BALL_NOISE + self.model.plate_noise = _DEFAULT_PLATE_NOISE + self.model.update_plate(plate_reset=True) + if self._use_dr: + self.model.set_initial_ball( + self._random_state.uniform( + low=-_PLATE_RADIUS * 0.75, + high=_PLATE_RADIUS * 0.75 + ), + self._random_state.uniform( + low=-_PLATE_RADIUS * 0.75, + high=_PLATE_RADIUS * 0.75 + ), + 0.0 + ) + self.model.ball_vel.x = self._random_state.uniform( + low=-0.02, high=0.02) + self.model.ball_vel.y = self._random_state.uniform( + low=-0.02, high=0.02) + else: + self.model.set_initial_ball( + -_PLATE_RADIUS * 0.5, + -_PLATE_RADIUS * 0.5, + 0.0 + ) + self.model.ball_vel.x = 0.0 + self.model.ball_vel.y = 0.0 + self.model.ball_vel.z = 0.0 + # self._set_velocity_for_speed_and_direction(0.0, 0.0) + self.model.iteration_count = 0 + return self._get_obs_from_state() + + def _get_obs_from_state(self) -> np.ndarray: + """ + Get the observation data from the simulator state. + """ + full_state = self.model.state() + # NOTE: The observation can be: + # estimated_x, estimated_y, + # estimated_velocity_x, estimated_velocity_y + if self._use_normalize_state: + c_ball_x = float(np.clip( + full_state["ball_x"], + -1 * _PLATE_RADIUS, + _PLATE_RADIUS + )) + c_ball_y = float(np.clip( + full_state["ball_y"], + -1 * _PLATE_RADIUS, + _PLATE_RADIUS + )) + c_ball_vel_x = float(np.clip( + full_state["ball_vel_x"], + -1 * _MAX_VELOCITY, + _MAX_VELOCITY + )) + c_ball_vel_y = float(np.clip( + full_state["ball_vel_y"], + -1 * _MAX_VELOCITY, + _MAX_VELOCITY + )) + return np.array( + [c_ball_x / (2 * _PLATE_RADIUS), + c_ball_y / (2 * _PLATE_RADIUS), + c_ball_vel_x / (2 * _MAX_VELOCITY), + c_ball_vel_y / (2 * _MAX_VELOCITY) + ] + ) + else: + return np.array( + [full_state["ball_x"], full_state["ball_y"], + full_state["ball_vel_x"], full_state["ball_vel_y"] + ] + ) + + def _get_reward(self) -> float: + """ + Get the reward value on a simulation step. + """ + # NOTE: The reward definition follows our inkling file + # moab.ink in brain/src/sdk2/samples/moabsim/moab.ink + full_state = self.model.state() + if self._is_penalty_terminal(): + return -10 + dx = full_state["ball_x"] - full_state["target_pos_x"] + dy = full_state["ball_y"] - full_state["target_pos_y"] + ball_origin_at_center = full_state["ball_radius"] + full_state[ + "plate_pos_z"] + _PLATE_Z_OFFSET + dz = full_state["ball_z"] - ball_origin_at_center + + distance_to_target = (dx ** 2 + dy ** 2 + dz ** 2) ** 0.5 + if distance_to_target < _CLOSE_ENOUGH: + return 10 + return 10 * _CLOSE_ENOUGH / distance_to_target + + def _is_penalty_terminal(self) -> bool: + """ + Check if it is the valid terminal for getting penalty. + """ + full_state = self.model.state() + return full_state["ball_fell_off"] > 0 or full_state[ + "iteration_count"] >= _MAX_ITER_COUNT + + def _set_velocity_for_speed_and_direction( + self, + speed: float, + direction: float + ) -> None: + """ + Set velocity for the ball speed and direction. + """ + # get the heading + dx = self.model.target_pos_x - self.model.ball.x + dy = self.model.target_pos_y - self.model.ball.y + + # direction is meaningless if we're already at the target + if (dx != 0) or (dy != 0): + + # set the magnitude + vel = vector.set_length([dx, dy, 0.0], speed) + + # rotate by direction around Z-axis at ball position + rot = matrix33.create_from_axis_rotation( + [0.0, 0.0, 1.0], + direction + ) + vel = matrix33.apply_to_vector(rot, vel) + + # unpack into ball velocity + self.model.ball_vel.x = vel[0] + self.model.ball_vel.y = vel[1] + self.model.ball_vel.z = vel[2] diff --git a/python/ray/rllib/tests/moab_env/moab_model.py b/python/ray/rllib/tests/moab_env/moab_model.py new file mode 100644 index 000000000000..3fbc11897429 --- /dev/null +++ b/python/ray/rllib/tests/moab_env/moab_model.py @@ -0,0 +1,773 @@ +""" +Simulator for the Moab plate+ball balancing device. +Author: Mike Estee +Copyright 2020 Microsoft +""" + +# pyright: strict + +import math +import random +import numpy as np +from typing import Dict, Tuple +import logging as log + +# NOTE: pyrr can create Integer vec/quats if you're not carefull! +# If you start seeing vectors zeroed out, make sure you're using float! +from pyrr import Quaternion, Vector3, matrix44, quaternion, vector, ray +from pyrr.geometric_tests import point_height_above_plane, ray_intersect_plane +from pyrr.plane import create_from_position + +# Some type aliases for clarity +Plane = np.ndarray +Ray = np.ndarray + +DEFAULT_SIMULATION_RATE = 0.020 # s, 20ms +DEFAULT_GRAVITY = 9.81 # m/s^2, Earth: there's no place like it. + +DEFAULT_BALL_RADIUS = 0.02 # m, Ping-Pong ball: 20mm +DEFAULT_BALL_SHELL = 0.0002 # m, Ping-Pong ball: 0.2mm +DEFAULT_BALL_MASS = 0.0027 # kg, Ping-Pong ball: 2.7g +DEFAULT_BALL_COR = 0.89 # unitless, Ping-Pong Coefficient Of Restitution: 0.89 +DEFAULT_FRICTION = 0.5 # unitless, Ping-Pong on Acrylic + +DEFAULT_PLATE_RADIUS = 0.225 / 2.0 # m, Moab: 225mm dia +PLATE_ORIGIN_TO_SURFACE_OFFSET = 0.009 # 9mm offset from plate rot origin to plate surface + +# plate limits +PLATE_HEIGHT_MAX = 0.040 # m, Moab: 40mm +DEFAULT_PLATE_HEIGHT = PLATE_HEIGHT_MAX / 2.0 +DEFAULT_PLATE_ANGLE_LIMIT = math.radians(44.0 * 0.5) # rad, 1/2 full range +DEFAULT_HEIGHT_Z_LIMIT = PLATE_HEIGHT_MAX / 2.0 # m, +/- limit from center Z pos + +# default ball Z position +DEFAULT_BALL_Z_POSITION = DEFAULT_PLATE_HEIGHT + PLATE_ORIGIN_TO_SURFACE_OFFSET + DEFAULT_BALL_RADIUS + +PLATE_MAX_Z_VELOCITY = 1.0 # m/s +PLATE_Z_ACCEL = 10.0 # m/s^2 + +# Moab measured velocity at 15deg in 3/60ths, or 300deg/s +DEFAULT_PLATE_MAX_ANGULAR_VELOCITY = (60.0 / 3.0) * math.radians(15) # rad/s + +# Set acceleration to get the plate up to velocity in 1/100th of a sec +DEFAULT_PLATE_ANGULAR_ACCEL = (100.0 / 1.0) * DEFAULT_PLATE_MAX_ANGULAR_VELOCITY # rad/s^2 + +# useful constants +X_AXIS = np.array([1.0, 0.0, 0.0]) +Y_AXIS = np.array([0.0, 1.0, 0.0]) +Z_AXIS = np.array([0.0, 0.0, 1.0]) + +# Sensor Actuator Noises +DEFAULT_PLATE_NOISE = 0.0 # radians +DEFAULT_BALL_NOISE = 0.0005 # add a little noise on perceived ball location in meters +DEFAULT_JITTER = 0.000 # No simulation jitter (s) + + +def clamp(val: float, min_val: float, max_val: float): + return min(max_val, max(min_val, val)) + + +# returns a noise value in the range [-scalar .. scalar] with a gaussian distribution +def random_noise(scalar: float) -> float: + return scalar * clamp(random.gauss(mu=0, sigma=0.333), -1, 1) # mean zero gauss with sigma = ~sqrt(scalar)/3 + + +class MoabModel: + def __init__(self): + self.pitch = 0.0 # type: float + self.roll = 0.0 # type: float + self.height_z = 0.0 # type: float + self.time_delta = 0.0 # type: float + self.jitter = 0.0 # type: float + self.step_time = 0.0 # type: float + self.elapsed_time = 0.0 # type: float + self.gravity = 0.0 # type: float + self.plate_radius = 0.0 # type: float + self.friction = 0.0 # type: float + self.tilt_max_vel = 0.0 # type: float + self.tilt_acc = 0.0 # type: float + self.ball_noise = 0.0 # type: float + self.plate_noise = 0.0 # type: float + self.ball_mass = 0.0 # type: float + self.ball_radius = 0.0 # type: float + self.ball_shell = 0.0 # type: float + self.ball_COR = 0.0 # type: float + self.plate_pos = Vector3() # type: Vector3 + + self.plate_theta_x = 0.0 # type: float + self.plate_theta_y = 0.0 # type: float + + self.plate_theta_vel_x = 0.0 # type: float + self.plate_theta_vel_y = 0.0 # type: float + self.plate_pos_vel_z = 0.0 # type: float + self.ball = Vector3() # type: Vector3 + self.ball_vel = Vector3() # type: Vector3 + self.ball_acc = Vector3() # type: Vector3 + self.ball_qat = Quaternion() # type: Quaternion + self.estimated_x = 0.0 # type: float + self.estimated_y = 0.0 # type: float + self.estimated_radius = 0.0 # type: float + self.estimated_vel_x = 0.0 # type: float + self.estimated_vel_y = 0.0 # type: float + self.prev_estimated_x = 0.0 # type: float + self.prev_estimated_y = 0.0 # type: float + self.iteration_count = 0 # type: int + self.experimental_physics = 0 # type: int + + self.reset() + + # set up all model parameters with sensable defaults + def reset(self): + # control inputs are Z-up, X-right, Y-forward + # control input (unitless) [-1..1] + self.pitch = 0.0 + self.roll = 0.0 + self.height_z = 0.0 + + # servo positions inputs, unitless [0..1] + # servos are PWM phase controlled. + # self.motors = [0, 0, 0] + + # model coordinate system is also Z-up + # units are Metric + # constants + self.time_delta = DEFAULT_SIMULATION_RATE + self.jitter = DEFAULT_JITTER + self.step_time = self.time_delta + self.elapsed_time = 0.0 + + self.gravity = DEFAULT_GRAVITY + self.plate_radius = DEFAULT_PLATE_RADIUS + self.friction = DEFAULT_FRICTION + self.tilt_max_vel = DEFAULT_PLATE_MAX_ANGULAR_VELOCITY + self.tilt_acc = DEFAULT_PLATE_ANGULAR_ACCEL + self.tilt_limit = DEFAULT_PLATE_ANGLE_LIMIT + self.height_z_limit = DEFAULT_HEIGHT_Z_LIMIT + self.ball_noise = DEFAULT_BALL_NOISE + self.plate_noise = DEFAULT_PLATE_NOISE + self.ball_mass = DEFAULT_BALL_MASS + self.ball_radius = DEFAULT_BALL_RADIUS + self.ball_shell = DEFAULT_BALL_SHELL + self.ball_COR = DEFAULT_BALL_COR + + # Coordinate System: + # - world coordinate space origin is the plate surface origin + # at it's lowest resting position. + # - plate origin is the Z-up top surface coordinate at + # the center of the plate. + # - ball position is world origin relative + + self.target_pos_x = 0 + self.target_pos_y = 0 + + # plate position (m) + # Y == plate height mid-point + self.plate_pos = Vector3([0.0, 0.0, DEFAULT_PLATE_HEIGHT]) + + # setting these sets the plate normal (rad) + self.plate_theta_x = 0.0 + self.plate_theta_y = 0.0 + + # control velocities for plate (rad/s) + self.plate_theta_vel_x = 0 + self.plate_theta_vel_y = 0 + self.plate_pos_vel_z = 0 # (m/s) + + # Linear position of ball (m) + self.ball = Vector3([ + 0.0, + 0.0, + # this should be equivalent to DEFAULT_BALL_Z_POSITION + self.plate_pos.z + PLATE_ORIGIN_TO_SURFACE_OFFSET + DEFAULT_BALL_RADIUS + ]) + + # Linear velocity of ball (m/sec) + self.ball_vel = Vector3([0.0, 0.0, 0.0]) + + # Linear accel of ball (m/sec^2) + self.ball_acc = Vector3([0.0, 0.0, 0.0]) + + # Orientation of ball (unitless) + self.ball_qat = Quaternion([0.0, 0.0, 0.0, 1.0]) # identity + + # Modelled camera observations + self.estimated_x = 0.0 + self.estimated_y = 0.0 + self.estimated_radius = self.ball_radius + + self.estimated_vel_x = 0.0 + self.estimated_vel_y = 0.0 + + # These are used to calculate estimated_vel_x/y from previous position + self.prev_estimated_x = 0.0 + self.prev_estimated_y = 0.0 + + # Meta variables + self.iteration_count = 0 # int [0..N] + self.experimental_physics = 0 + + def halted(self) -> bool: + # ball is still on the plate? + return self.get_ball_distance_to_center() > self.plate_radius + + def step(self): + # update the step time for this step + self.step_time = self.time_delta + random_noise(self.jitter) + self.elapsed_time += self.step_time + + # update plate metrics from inputs + self.update_plate() + + # update ball metrics + self.update_ball() + + # update meta + self.iteration_count += 1 + + # single axis acceleration of a pos towards a destination + # with a hard stop at the destination + def _accel_param(self, + q: float, dest: float, vel: float, acc: float, max_vel: float): + + # direction of accel + dir = 0.0 + if q < dest: + dir = 1.0 + if q > dest: + dir = -1.0 + + # calculate the change in velocity and position + acc = acc * dir * self.step_time + vel_end = clamp(vel + acc * self.step_time, -max_vel, max_vel) + vel_avg = (vel + vel_end) * 0.5 + delta = vel_avg * self.step_time + vel = vel_end + + # moving towards the dest? + if (dir > 0 and q < dest and q + delta < dest) or \ + (dir < 0 and q > dest and q + delta > dest): + q = q + delta + + # stop at dest + else: + q = dest + vel = 0 + + return (q, vel) + + # convert X/Y theta components into a Z-Up RH plane normal + def _nor_from_xy_theta(self, x_theta: float, y_theta: float) -> np.ndarray: + x_rot = matrix44.create_from_axis_rotation(axis=X_AXIS, theta=x_theta) + y_rot = matrix44.create_from_axis_rotation(axis=Y_AXIS, theta=y_theta) + + # pitch then roll + nor = matrix44.apply_to_vector(mat=x_rot, vec=Z_AXIS) + nor = matrix44.apply_to_vector(mat=y_rot, vec=nor) + + nor = vector.normalize(nor) + return nor + + def _plate_nor(self) -> Vector3: + return Vector3(self._nor_from_xy_theta(self.plate_theta_x, self.plate_theta_y)) + + def update_plate(self, plate_reset: bool = False): + # Find the target xth,yth & zpos + # convert xy[-1..1] to zx[-self.tilt_limit .. self.tilt_limit] + # convert z[-1..1] to [PLATE_HEIGHT_MAX/2 - self.height_z_limit .. PLATE_HEIGHT_MAX/2 + self.height_z_limit] + theta_x_target = self.tilt_limit * self.pitch # pitch around X axis + theta_y_target = self.tilt_limit * self.roll # roll around Y axis + z_target = (self.height_z * self.height_z_limit) + PLATE_HEIGHT_MAX / 2.0 + + # quantize target positions to whole degree increments + # the Moab hardware can only command by whole degrees + theta_y_target = math.radians(round(math.degrees(theta_y_target))) + theta_x_target = math.radians(round(math.degrees(theta_x_target))) + + # get the current xth,yth & zpos + theta_x, theta_y = self.plate_theta_x, self.plate_theta_y + z_pos = self.plate_pos.z + + # on reset, bypass the motion equations + if plate_reset: + theta_x = theta_x_target + theta_y = theta_y_target + z_pos = z_target + + # smooth transition to target based on accel and velocity limits + else: + theta_x, self.plate_theta_vel_x = \ + self._accel_param(theta_x, theta_x_target, self.plate_theta_vel_x, self.tilt_acc, self.tilt_max_vel) + theta_y, self.plate_theta_vel_y = \ + self._accel_param(theta_y, theta_y_target, self.plate_theta_vel_y, self.tilt_acc, self.tilt_max_vel) + z_pos, self.plate_pos_vel_z = \ + self._accel_param(z_pos, z_target, self.plate_pos_vel_z, PLATE_Z_ACCEL, PLATE_MAX_Z_VELOCITY) + + # add noise to the plate positions + theta_x += random_noise(self.plate_noise) + theta_y += random_noise(self.plate_noise) + + # clamp to range limits + theta_x = clamp(theta_x, -self.tilt_limit, self.tilt_limit) + theta_y = clamp(theta_y, -self.tilt_limit, self.tilt_limit) + z_pos = clamp(z_pos, + PLATE_HEIGHT_MAX / 2.0 - self.height_z_limit, + PLATE_HEIGHT_MAX / 2.0 + self.height_z_limit) + + # Now convert back to plane parameters + self.plate_theta_x = theta_x + self.plate_theta_y = theta_y + self.plate_pos.z = z_pos + + # ball intertia with radius and hollow radius + # I = 2/5 * m * ((r^5 - h^5) / (r^3 - h^3)) + def _ball_inertia(self): + hollow_radius = self.ball_radius - self.ball_shell + return 2.0 / 5.0 * self.ball_mass * ( + (math.pow(self.ball_radius, 5.0) - math.pow(hollow_radius, 5.0)) / + (math.pow(self.ball_radius, 3.0) - math.pow(hollow_radius, 3.0))) + + def _camera_pos(self) -> Vector3: + """ camera origin (lens center) in world space """ + return Vector3([0., 0., -0.052]) + + def _estimated_ball(self, ball: Vector3) -> Tuple[float, float, float]: + """ + ray trace the ball position and an edge of the ball back to the camera + origin and use the collision points with the tilted plate to estimate + what a camera might perceive the ball position and size to be. + """ + # contact ray from camera to plate + camera = self._camera_pos() + displacement = camera - self.ball + displacement_radius = camera - (self.ball + Vector3([self.ball_radius, 0, 0])) + + ball_ray = ray.create(camera, displacement) + ball_radius_ray = ray.create(camera, displacement_radius) + + surface_plane = self._surface_plane() + + contact = Vector3(ray_intersect_plane(ball_ray, surface_plane, False)) + radius_contact = Vector3(ray_intersect_plane(ball_radius_ray, surface_plane, False)) + + x, y = contact.x, contact.y + r = math.fabs(contact.x - radius_contact.x) + + # add the noise in + estimated_x = x + random_noise(self.ball_noise) + estimated_y = y + random_noise(self.ball_noise) + estimated_radius = r + random_noise(self.ball_noise) + + return estimated_x, estimated_y, estimated_radius + + def _estimated_speed(self) -> float: + return vector.length([self.ball_vel.x, self.ball_vel.y, self.ball_vel.z]) + + def _estimated_direction(self) -> float: + # get the vector to the target + dx = self.target_pos_x - self.estimated_x + dy = self.target_pos_y - self.estimated_y + + # vectors and lengths + u = vector.normalize([dx, dy, 0.0]) + v = vector.normalize([self.estimated_vel_x, self.estimated_vel_y, 0.0]) + ul = vector.length(u) + vl = vector.length(v) + + # no velocity? already on the target? + if ul == 0.0 or vl == 0.0: + return 0.0 + else: + # angle between vectors + uv_dot = vector.dot(u, v) + + # signed angle + x = u[0] + y = u[1] + angle = math.atan2(vector.dot([-y, x, 0.0], v), uv_dot) + if math.isnan(angle): + return 0.0 + else: + return angle + + def _surface_plane(self) -> Plane: + """ + Return the surface plane of the plate + """ + plate_surface = np.array([ + self.plate_pos.x, + self.plate_pos.y, + self.plate_pos.z + PLATE_ORIGIN_TO_SURFACE_OFFSET + ]) + return create_from_position(plate_surface, self._plate_nor()) + + def _ball_rest_plane(self) -> Plane: + """ + Return the plane which represents all the valid positions + for the ball origin at rest. + """ + plate_surface = np.array([ + self.plate_pos.x, + self.plate_pos.y, + self.plate_pos.z + PLATE_ORIGIN_TO_SURFACE_OFFSET + self.ball_radius + ]) + return create_from_position(plate_surface, self._plate_nor()) + + def _motion_for_time(self, u: Vector3, a: Vector3, t: float) -> Tuple[Vector3, Vector3]: + """ + Equations of motion for displacement and final velocity + u: initial velocity + a: acceleration + d: displacement + v: final velocity + + d = ut + 1/2at^2 + v = u + at + + returns (d, v) + """ + d = (u * t) + (0.5 * a * (t ** 2)) + v = u + a * t + return d, v + + def _time_for_motion(self, u: Vector3, a: Vector3, d: Vector3) -> float: + """ + Compute the time delta given: + u: initial velocity + a: acceleration + d: displacement + + t = iff a>0 and 2*d*a+u**2>=0; (-u + sqrt(2*d*a + u**2))/a + t = iff a<0 and 2*d*a+u**2>=0; -(u + sqrt(2*d*a + u**2))/a + t = iff a==0; 0 + """ + + # here's the t = (v - u) / a form: + # if a.x != 0.0: + # t = (v.x - u.x) / a.x + # elif a.y != 0.0: + # t = (v.y - u.y) / a.y + # elif a.z != 0.0: + # t = (v.z - u.z) / a.z + + t = 0.0 + k = 2 * d.x * a.x + (u.x ** 2) + if a.x > 0.0 and k > 0.0: + t = (-u.x + math.sqrt(k)) / a.x + elif a.x < 0.0 and k > 0.0: + t = -(u.x + math.sqrt(k)) / a.x + + k = 2 * d.y * a.y + (u.y ** 2) + if a.y > 0.0 and k > 0.0: + t = (-u.y + math.sqrt(k)) / a.y + elif a.y < 0.0 and k > 0.0: + t = -(u.y + math.sqrt(k)) / a.y + + k = 2 * d.z * a.z + (u.z ** 2) + if a.z > 0.0 and k > 0.0: + t = (-u.z + math.sqrt(k)) / a.z + elif a.z < 0.0 and k > 0.0: + t = -(u.z + math.sqrt(k)) / a.z + + return t + + def _ball_airbone(self, step_t: float) -> float: + """ + Calculate the final position of the ball while airborne. + + if a collision happens: + - updates the ball position and velocity at point of contact + - returns the remaining time in the simulation + """ + + # calculate the acceleration + self.ball_acc = Vector3([0.0, 0.0, -self.gravity]) + self.ball_rot_acc = Vector3([0.0, 0.0, 0.0]) # no change in a vacuum + + # calculate the displacement, new velocity, and remaining time for this simulation step + dispplacement, vel = self._motion_for_time(self.ball_vel, self.ball_acc, step_t) + remaining_t = 0.0 + + # get the collision plane, which is the contact plane + offset by the ball radius + ball_rest = self._ball_rest_plane() + collision = point_height_above_plane(self.ball + dispplacement, ball_rest) < 0 + if collision: + # a ray intersecting a plane + motion_ray = ray.create(self.ball, dispplacement) + ball_rest = self._ball_rest_plane() + contact = Vector3(ray_intersect_plane(motion_ray, ball_rest, False)) + + # invert equations of motion, solve for time at contact pt + contact_disp = contact - self.ball + contact_t = self._time_for_motion(self.ball_vel, self.ball_acc, contact_disp) + remaining_t = step_t - contact_t + + # re-calc motion at time + dispplacement, vel = self._motion_for_time(self.ball_vel, self.ball_acc, contact_t) + + # update ball position and velocities + self.ball.x += dispplacement.x + self.ball.y += dispplacement.y + self.ball.z += dispplacement.z + self.ball_vel = vel + + # return the remaining time in the equation, if any + # this implicitly means that the ball has collided with the plate + # and we will transition to the collision state + return remaining_t + + def _ball_plate_collision(self, remaining_t: float) -> float: + # Using the velocity as a incident ray, reflect it off the + # ball plate for the new velocity at the contact point. + # Where: + # R: reflected vector + # I: incident vector + # N: normal of plate + # R = I - 2N(I dot N) + vel = self.ball_vel + nor = self._plate_nor() + vel_refl = vel - 2.0 * nor * (vel | nor) + + # Apply KE loss using Coefficient Of Restitution. + # e: Coefficient Of Restitution + # KE: kinetic energy + # m: mass + # u: pre-collision velocity + # v: post-collision velocity + # KE = 0.5 * m * v^2 + # e = sqrt(0.5 * m * v^2 / 0.5 * m * u^2) + # v = u*e + vel_len = vector.length(vel_refl) * self.ball_COR + vel_final = Vector3(vector.set_length(vel_refl, vel_len)) + + # set the reflected velocity + self.ball_vel = vel_final + return remaining_t + + def _update_ball_z(self): + self.ball.z = ( + self.ball.x * math.sin(-self.plate_theta_y) + + self.ball.y * math.sin(self.plate_theta_x) + + self.ball_radius + + self.plate_pos.z + + PLATE_ORIGIN_TO_SURFACE_OFFSET + ) + + def _ball_plate_contact(self, step_t: float) -> float: + # NOTE: the x_theta axis creates motion in the Y-axis, and vice versa + # x_theta, y_theta = self._xy_theta_from_nor(self.plate_nor.xyz) + x_theta = self.plate_theta_x + y_theta = self.plate_theta_y + + # Equations for acceleration on a plate at rest + # accel = (mass * g * theta) / (mass + inertia / radius^2) + # (y_theta,x are intentional swapped here.) + theta = Vector3([y_theta, -x_theta, 0]) + self.ball_acc = ( + theta / (self.ball_mass + self._ball_inertia() / (self.ball_radius ** 2)) * + self.ball_mass * self.gravity + ) + + # get contact displacement + disp, vel = self._motion_for_time(self.ball_vel, self.ball_acc, step_t) + + # simplified ball mechanics against a plane + self.ball.x += disp.x + self.ball.y += disp.y + self._update_ball_z() + self.ball_vel = vel + + # For rotation on plate motion we use infinite friction and + # perfect ball / plate coupling. + # Calculate the distance we traveled across the plate during + # this time slice. + rot_distance = math.hypot(disp.x, disp.y) + if rot_distance > 0: + # Calculate the fraction of the circumference that we traveled + # (in radians). + rot_angle = rot_distance / self.ball_radius + + # Create a quaternion that represents the delta rotation for this time period. + # Note that we translate the (x, y) direction into (y, -x) because we're + # creating a vector that represents the axis of rotation which is normal + # to the direction the ball traveled in the x/y plane. + rot_q = quaternion.normalize( + np.array([ + disp.y / rot_distance * math.sin(rot_angle / 2.0), + -disp.x / rot_distance * math.sin(rot_angle / 2.0), + 0.0, + math.cos(rot_angle / 2.0) + ]) + ) + + old_rot = self.ball_qat.xyzw + new_rot = quaternion.cross(quat1=old_rot, quat2=rot_q) + self.ball_qat.xyzw = quaternion.normalize(new_rot) + return 0.0 + + def _airborne(self) -> bool: + """ + returns: True when the ball is not in contact with the plate + """ + contact_plane = self._ball_rest_plane() + dist_to_contact = point_height_above_plane(self.ball, contact_plane) + + if dist_to_contact <= 0.0: + return False + else: + return True + + def set_initial_ball(self, x: float, y: float, z: float): + self.ball.xyz = [x, y, z] + self._update_ball_z() + + # Set initial observations + self.estimated_x, self.estimated_y, self.estimated_radius = self._estimated_ball(self.ball) + + self.prev_estimated_x = self.estimated_x + self.prev_estimated_y = self.estimated_y + pass + + def update_ball(self): + """ + Update the ball position with the physics model. We have three phases: + - airborn: ball is not contacting the plate + - collision: the instant the ball contacts the plate + - contact: the ball is contacting the plate + """ + remaining_t = self.step_time + + # experimental bounce physics + if self.experimental_physics != 0: + if self._airborne(): + log.debug("airborne") + remaining_t = self._ball_airbone(remaining_t) + if remaining_t > 0.0: + log.debug("collision") + remaining_t = self._ball_plate_collision(remaining_t) + remaining_t = self._ball_airbone(remaining_t) + + # ESTEE: swallow small bounces that would re-collide with the plate... + else: + log.debug("contact") + remaining_t = self._ball_plate_contact(self.step_time) + + # contact only... + else: + remaining_t = self._ball_plate_contact(self.step_time) + + # Finally, lets make some approximations for observations + self.estimated_x, self.estimated_y, self.estimated_radius = self._estimated_ball(self.ball) + + # Use n-1 states to calculate an estimated velocity. + self.estimated_vel_x = (self.estimated_x - self.prev_estimated_x) / self.step_time + self.estimated_vel_y = (self.estimated_y - self.prev_estimated_y) / self.step_time + + # update for next time + self.prev_estimated_x = self.estimated_x + self.prev_estimated_y = self.estimated_y + pass + + def get_ball_distance_to_center(self) -> float: + # ball.z relative to plate + zpos = self.ball.z - (self.plate_pos.z + self.ball_radius + + PLATE_ORIGIN_TO_SURFACE_OFFSET) + + # ball distance from ball position on plate at origin + return math.sqrt( + math.pow(self.ball.x, 2.0) + + math.pow(self.ball.y, 2.0) + + math.pow(zpos, 2.0)) + + def get_ball_velocity(self) -> float: + return vector.length(self.ball_vel.xyz) + + def state(self) -> Dict[str, float]: + # x_theta, y_theta = self._xy_theta_from_nor(self.plate_nor) + plate_nor = self._plate_nor() + + return dict( + + # reflected input controls + tilt_x=self.roll, + tilt_y=self.pitch, + + roll=self.roll, + pitch=self.pitch, + + height_z=self.height_z, + + # reflected constants + time_delta=self.time_delta, + jitter=self.jitter, + step_time=self.step_time, + elapsed_time=self.elapsed_time, + + gravity=self.gravity, + plate_radius=self.plate_radius, + friction=self.friction, + tilt_max_vel=self.tilt_max_vel, + tilt_acc=self.tilt_acc, + tilt_limit=self.tilt_limit, + height_z_limit=self.height_z_limit, + + ball_mass=self.ball_mass, + ball_radius=self.ball_radius, + ball_shell=self.ball_shell, + ball_COR=self.ball_COR, + + target_pos_x=self.target_pos_x, + target_pos_y=self.target_pos_y, + + # modelled plate metrics + plate_pos_x=self.plate_pos.x, + plate_pos_y=self.plate_pos.y, + plate_pos_z=self.plate_pos.z, + + plate_nor_x=plate_nor.x, + plate_nor_y=plate_nor.y, + plate_nor_z=plate_nor.z, + + plate_theta_x=self.plate_theta_x, + plate_theta_y=self.plate_theta_y, + + plate_theta_vel_x=self.plate_theta_vel_x, + plate_theta_vel_y=self.plate_theta_vel_y, + plate_pos_vel_z=self.plate_pos_vel_z, + + # modelled ball metrics + ball_x=self.ball.x, + ball_y=self.ball.y, + ball_z=self.ball.z, + + ball_vel_x=self.ball_vel.x, + ball_vel_y=self.ball_vel.y, + ball_vel_z=self.ball_vel.z, + + ball_acc_x=self.ball_acc.x, + ball_acc_y=self.ball_acc.y, + ball_acc_z=self.ball_acc.z, + + ball_qat_x=self.ball_qat.x, + ball_qat_y=self.ball_qat.y, + ball_qat_z=self.ball_qat.z, + ball_qat_w=self.ball_qat.w, + + # modelled camera observations + estimated_x=self.estimated_x, + estimated_y=self.estimated_y, + estimated_radius=self.estimated_radius, + + estimated_vel_x=self.estimated_vel_x, + estimated_vel_y=self.estimated_vel_y, + + estimated_speed=self._estimated_speed(), + estimated_direction=self._estimated_direction(), + + ball_noise=self.ball_noise, + plate_noise=self.plate_noise, + + # meta vars + ball_fell_off=1 if self.halted() else 0, + iteration_count=self.iteration_count, + experimental_physics=self.experimental_physics + ) diff --git a/python/ray/rllib/tests/test_convergence.py b/python/ray/rllib/tests/test_convergence.py index e2e6a6778d0b..f96fc26710db 100644 --- a/python/ray/rllib/tests/test_convergence.py +++ b/python/ray/rllib/tests/test_convergence.py @@ -4,6 +4,7 @@ # from ray.rllib.agents.ppo import DEFAULT_CONFIG, PPOTrainer as Trainer from ray.rllib.agents.sac import DEFAULT_CONFIG, SACTrainer as Trainer +from moab_env.moab_env import MoabSim class GymEnv: @@ -26,6 +27,7 @@ def ray_env(): ray.shutdown() +""" @pytest.fixture def config(env_id): config = DEFAULT_CONFIG.copy() @@ -35,23 +37,63 @@ def config(env_id): config["num_workers"] = 8 config["num_envs_per_worker"] = 8 return config +""" @pytest.fixture def trainer(config): - return Trainer(config=config, env=GymEnv) + return Trainer(config=config, env=MoabSim) + + +#@pytest.mark.parametrize( +# "env_id, threshold", [ +# ("Pendulum-v0", -750), +# ] +#) +#@pytest.mark.usefixtures("ray_env") +#def test_convergence(trainer, env_id, threshold): +# mean_episode_reward = -float("inf") +# for i in range(500): +# result = trainer.train() +# mean_episode_reward = result["episode_reward_mean"] +# print(f"{i}: {mean_episode_reward}") +# assert mean_episode_reward >= threshold + + +@pytest.fixture +def config(env_id): + config = DEFAULT_CONFIG.copy() + config["env_config"] = { + "use_normalize_action": False, + "use_dr": False, + "use_normalize_state": True + } + config["evaluation_config"] = { + "exploration_enabled": False + } + config["evaluation_interval"] = 1 + config["evaluation_num_episodes"] = 10 + if env_id == "Moab-v0": + config["timesteps_per_iteration"] = 2000 + config["num_workers"] = 8 + config["num_envs_per_worker"] = 8 + config["learning_starts"] = 6000 + config["normalize_actions"] = True + return config @pytest.mark.parametrize( "env_id, threshold", [ - ("Pendulum-v0", -750), + ("Moab-v0", 1000), ] ) @pytest.mark.usefixtures("ray_env") def test_convergence(trainer, env_id, threshold): mean_episode_reward = -float("inf") - for i in range(500): + for i in range(10000): result = trainer.train() mean_episode_reward = result["episode_reward_mean"] - print(f"{i}: {mean_episode_reward}") - assert mean_episode_reward >= threshold + total_samples = result["timesteps_total"] + print(f"SAC Iteration {i}: Train Episode Reward Mean: " + f"{mean_episode_reward}, " + f"Number of samples {total_samples}")