Skip to content

Commit

Permalink
most things passing
Browse files Browse the repository at this point in the history
  • Loading branch information
NishanthJKumar committed Oct 26, 2023
1 parent c7d9736 commit cebc2be
Show file tree
Hide file tree
Showing 6 changed files with 70 additions and 81 deletions.
3 changes: 1 addition & 2 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -27,5 +27,4 @@ tests/_fake_results
sas_plan

# Jetbrains IDEs
.idea/
analyze_sticky_data.py
.idea/
87 changes: 56 additions & 31 deletions predicators/envs/ball_and_cup_sticky_table.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,13 @@ class BallAndCupStickyTableEnv(BaseEnv):
tables. This environment is a more-complex (but significantly different)
version of the sticky-table environment.
Most of the tables are completely flat, but one is half is smooth and half
is sticky. If the agent tries to place the ball directly on any table,
it will roll off with high probability. If it tries to place it on a
half-flat half-smooth table, the ball will *certainly* roll off. However,
if the ball is placed inside a cup first, then the ball + cup system stays
on the sticky and normal table surfaces with high probability.
Most of the tables are completely flat, but one is half is mostly smooth
sticky in a particular circular region on the table. If the agent tries
to place the ball directly on any table, it will roll off with high
probability. If it tries to place it on a special table,
the ball will *certainly* roll off. However, if the ball is placed inside
a cup first, then the ball + cup system stays on the sticky and normal
table surfaces with high probability.
Note that unlike almost all of our other environments, there is real
stochasticity in the outcomes of placing.
Expand Down Expand Up @@ -148,7 +149,7 @@ def _place_smooth_fall_prob(self) -> float:
return CFG.sticky_table_place_smooth_fall_prob

@classmethod
def _object_to_geom(self, obj: Object, state: State) -> utils._Geom2D:
def _object_to_geom(cls, obj: Object, state: State) -> utils._Geom2D:
x = state.get(obj, "x")
y = state.get(obj, "y")
radius = state.get(obj, "radius")
Expand Down Expand Up @@ -182,7 +183,7 @@ def _get_tasks(self, num: int,
# Add a random spin to offset the circle. This is to ensure
# the tables are in different positions along the circle every
# time.
theta_offset = 0.0 #rng.uniform(0, 2 * np.pi)
theta_offset = 0.0 #rng.uniform(0, 2 * np.pi)
sticky_region_radius = radius * self.sticky_region_radius_scale
# Now, actually instantiate the tables.
for i, theta in enumerate(thetas):
Expand All @@ -195,16 +196,26 @@ def _get_tasks(self, num: int,
prefix = "sticky"
sticky = 1.0
obj = Object(f"{prefix}-table-{i}", self._table_type)
sticky_region_dist_from_center = rng.uniform(0.0, radius - sticky_region_radius)
sticky_region_dist_from_center = rng.uniform(
0.0, radius - sticky_region_radius)
sticky_region_theta_from_center = rng.uniform(0.0, 2 * np.pi)
state_dict[obj] = {
"x": x,
"y": y,
"radius": radius,
"sticky": sticky,
"sticky_region_x_offset": sticky_region_dist_from_center * np.cos(sticky_region_theta_from_center),
"sticky_region_y_offset": sticky_region_dist_from_center * np.sin(sticky_region_theta_from_center),
"sticky_region_radius": sticky_region_radius
"x":
x,
"y":
y,
"radius":
radius,
"sticky":
sticky,
"sticky_region_x_offset":
sticky_region_dist_from_center *
np.cos(sticky_region_theta_from_center),
"sticky_region_y_offset":
sticky_region_dist_from_center *
np.sin(sticky_region_theta_from_center),
"sticky_region_radius":
sticky_region_radius
}
tables = sorted(state_dict)
target_table = tables[-1]
Expand Down Expand Up @@ -270,15 +281,15 @@ def _get_tasks(self, num: int,
return tasks

@classmethod
def exists_robot_collision(self, state: State) -> bool:
def exists_robot_collision(cls, state: State) -> bool:
"""Return true if there is a collision between the robot and any other
object in the environment."""
robot, = state.get_objects(self._robot_type)
robot, = state.get_objects(cls._robot_type)
all_possible_collision_objs = state.get_objects(
self._table_type) + state.get_objects(
self._cup_type) + state.get_objects(self._ball_type)
cls._table_type) + state.get_objects(
cls._cup_type) + state.get_objects(cls._ball_type)
for obj in all_possible_collision_objs:
obj_geom = self._object_to_geom(obj, state)
obj_geom = cls._object_to_geom(obj, state)
if obj_geom.contains_point(state.get(robot, "x"),
state.get(robot, "y")):
return True
Expand Down Expand Up @@ -410,6 +421,7 @@ def simulate(self, state: State, action: Action) -> State:
# Placing logic.
else:
if not hand_empty:
assert obj_being_held is not None
# Find the table for placing, if any.
table: Optional[Object] = None
for target in state.get_objects(self._table_type):
Expand Down Expand Up @@ -441,26 +453,38 @@ def simulate(self, state: State, action: Action) -> State:
if obj_being_held is not None:
next_state.set(obj_being_held, "held", 0.0)
if obj_type_id == 3.0:
# Possibly put on the table, or have it fall somewhere near.
# Possibly put on the table, or have it fall
# somewhere near.
fall_prob = self._place_sticky_fall_prob
if obj_being_held == ball:
fall_prob = self._place_ball_fall_prob
if self._table_is_sticky(table, state):
# Check if placing on the smooth side of the sticky table,
# and set fall prob accordingly.
sticky_region_x = state.get(table, "sticky_region_x_offset") + table_x
sticky_region_y = state.get(table, "sticky_region_y_offset") + table_y
sticky_region = utils.Circle(sticky_region_x, sticky_region_y, state.get(table, "sticky_region_radius"))
# Check if placing on the smooth part of
# the sticky table, and set fall prob
# accordingly.
sticky_region_x = state.get(
table,
"sticky_region_x_offset") + table_x
sticky_region_y = state.get(
table,
"sticky_region_y_offset") + table_y
sticky_region = utils.Circle(
sticky_region_x, sticky_region_y,
state.get(table,
"sticky_region_radius"))
if not sticky_region.contains_point(
act_x, act_y):
if obj_being_held == cup:
fall_prob = self._place_smooth_fall_prob
fall_prob = \
self._place_smooth_fall_prob
else:
assert obj_being_held == ball
fall_prob = 1.0
# Handle object falling or placing on table surface.
# Handle object falling or placing on table
# surface.
if self._noise_rng.uniform() < fall_prob:
fall_x, fall_y = self._sample_floor_point_around_table(
fall_x, fall_y = \
self._sample_floor_point_around_table(
table, state, self._noise_rng)
next_state = self._handle_placing_object(
fall_x, fall_y, next_state,
Expand All @@ -476,7 +500,8 @@ def simulate(self, state: State, action: Action) -> State:
assert self._OnTable_holds(
next_state, [obj_being_held, table])
else:
assert obj_type_id == 2.0 # corresponding to placing in cup
# corresponding to placing in cup
assert obj_type_id == 2.0
assert obj_being_held == ball
next_state.set(ball, "x", act_x)
next_state.set(ball, "y", act_y)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -344,7 +344,7 @@ def place_ball_on_floor_sampler(state: State, goal: Set[GroundAtom],
def place_ball_in_cup_sampler(state: State, goal: Set[GroundAtom],
rng: np.random.Generator,
objs: Sequence[Object]) -> Array:
del goal # unused
del rng, goal # unused
cup = objs[2]
# Just place the ball in the middle of the cup. Set
# the type id to be 2.0 to correspond to the cup
Expand Down Expand Up @@ -427,10 +427,10 @@ def place_ball_in_cup_sampler(state: State, goal: Set[GroundAtom],
# LiftedAtom(HoldingBall, [ball])
# }
# placecupwithballontable_nsrt = NSRT("PlaceCupWithBallOnTable",
# parameters, preconditions,
# add_effects, delete_effects, set(),
# option, option_vars,
# place_on_table_sampler)
# parameters, preconditions,
# add_effects, delete_effects, set(),
# option, option_vars,
# place_on_table_sampler)
# nsrts.add(placecupwithballontable_nsrt)

# PlaceCupWithoutBallOnFloor
Expand Down Expand Up @@ -482,7 +482,7 @@ def place_ball_in_cup_sampler(state: State, goal: Set[GroundAtom],
option_vars = parameters
option = NavigateToBall
preconditions = set()
add_effects = {ReachableBall([robot, ball])}
add_effects = {LiftedAtom(ReachableBall, [robot, ball])}
ignore_effects = {ReachableSurface, ReachableBall, ReachableCup}

def navigate_to_obj_sampler(state: State, goal: Set[GroundAtom],
Expand Down Expand Up @@ -539,7 +539,7 @@ def navigate_to_obj_sampler(state: State, goal: Set[GroundAtom],
option_vars = parameters
option = NavigateToCup
preconditions = set()
add_effects = {ReachableCup([robot, cup])}
add_effects = {LiftedAtom(ReachableSurface, [robot, table])}
ignore_effects = {ReachableSurface, ReachableBall, ReachableCup}
navigatetocup_nsrt = NSRT("NavigateToCup", parameters, preconditions,
add_effects, set(), ignore_effects, option,
Expand All @@ -551,7 +551,7 @@ def navigate_to_obj_sampler(state: State, goal: Set[GroundAtom],
option_vars = parameters
option = NavigateToTable
preconditions = set()
add_effects = {ReachableSurface([robot, table])}
add_effects = {LiftedAtom(ReachableSurface, [robot, table])}
ignore_effects = {ReachableSurface, ReachableBall, ReachableCup}
navigatetotable_nsrt = NSRT("NavigateToTable",
parameters, preconditions, add_effects,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,8 @@ def get_options(cls, env_name: str, types: Dict[str, Type],
cup_type = types["cup"]
ball_type = types["ball"]
table_type = types["table"]
# Parameters are move_or_pickplace, obj_type_id, ball_only, absolute x, y actions.
# Parameters are move_or_pickplace, obj_type_id, ball_only,
# absolute x, y actions.
params_space = Box(
np.array([
0.0, 0.0, 0.0, BallAndCupStickyTableEnv.x_lb,
Expand Down
41 changes: 3 additions & 38 deletions predicators/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -301,42 +301,24 @@ def construct_active_sampler_input(state: State, objects: Sequence[Object],
sampler_input_lst.append(params[0] - target_pos)
elif CFG.env == "ball_and_cup_sticky_table":
if "PlaceCup" in param_option.name and "Table" in param_option.name:
robot, ball, cup, table = objects
robot_y = state.get(robot, "y")
robot_x = state.get(robot, "x")
_, _, _, table = objects
table_y = state.get(table, "y")
table_x = state.get(table, "x")
ball_x = state.get(ball, "x")
ball_y = state.get(ball, "y")
cup_x = state.get(cup, "x")
cup_y = state.get(cup, "y")
sticky = state.get(table, "sticky")
# sticky_radius = state.get(table, "sticky_radius")
sticky_region_x = state.get(table, "sticky_region_x_offset")
sticky_region_y = state.get(table, "sticky_region_y_offset")
sticky_region_radius = state.get(table, "sticky_region_radius")
table_radius = state.get(table, "radius")
a, b, c, param_x, param_y = params
_, _, _, param_x, param_y = params
sampler_input_lst.append(table_radius)
sampler_input_lst.append(sticky)
sampler_input_lst.append(sticky_region_x)
sampler_input_lst.append(sticky_region_y)
sampler_input_lst.append(sticky_region_radius)
# sampler_input_lst.append(ball_x)
# sampler_input_lst.append(ball_y)
# sampler_input_lst.append(cup_x)
# sampler_input_lst.append(cup_y)
# sampler_input_lst.append(robot_x)
# sampler_input_lst.append(robot_y)
sampler_input_lst.append(table_x)
sampler_input_lst.append(table_y)
# sampler_input_lst.append(a)
# sampler_input_lst.append(b)
# sampler_input_lst.append(c)
sampler_input_lst.append(param_x)
sampler_input_lst.append(param_y)
# sampler_input_lst.append(param_x - table_x)
# sampler_input_lst.append(param_y - table_y)
else:
raise NotImplementedError("Oracle feature selection not "
f"implemented for {CFG.env}")
Expand Down Expand Up @@ -403,25 +385,8 @@ def plot(self, ax: plt.Axes, **kwargs: Any) -> None:
def contains_point(self, x: float, y: float) -> bool:
return (x - self.x)**2 + (y - self.y)**2 <= self.radius**2

def sector_contains_point(self, x: float, y: float,
sector_start_angle: float,
sector_end_angle: float) -> bool:
"""Returns true if the point x, y is contained within the sector
starting at sector_start_angle radians and ending at sector_end_angle
radians."""
# First, check that the point is even on the circle.
if not self.contains_point(x, y):
return False
# Next, convert (x, y) relative to the table's center
# to polar coordinates.
relative_x = x - self.x
relative_y = y - self.y
theta = np.arctan2(relative_y, relative_x)
if theta < 0:
theta = np.pi - theta
return sector_start_angle <= theta <= sector_end_angle

def contains_circle(self, other_circle: Circle) -> bool:
"""Check whether this circle wholly contains another one."""
dist_between_centers = np.sqrt((other_circle.x - self.x)**2 +
(other_circle.y - self.y)**2)
return (dist_between_centers + other_circle.radius) <= self.radius
Expand Down
1 change: 0 additions & 1 deletion tests/envs/test_ball_and_cup_sticky_table_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
from predicators import utils
from predicators.envs.ball_and_cup_sticky_table import BallAndCupStickyTableEnv
from predicators.ground_truth_models import get_gt_nsrts, get_gt_options
from predicators.structs import Action


def test_sticky_table():
Expand Down

0 comments on commit cebc2be

Please sign in to comment.